github.com/k14s/starlark-go@v0.0.0-20200720175618-3a5c849cc368/internal/compile/codegen_test.go (about)

     1  package compile
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"testing"
     7  
     8  	"github.com/k14s/starlark-go/resolve"
     9  	"github.com/k14s/starlark-go/syntax"
    10  )
    11  
    12  // TestPlusFolding ensures that the compiler generates optimized code for
    13  // n-ary addition of strings, lists, and tuples.
    14  func TestPlusFolding(t *testing.T) {
    15  	isPredeclared := func(name string) bool { return name == "x" }
    16  	isUniversal := func(name string) bool { return false }
    17  	for i, test := range []struct {
    18  		src  string // source expression
    19  		want string // disassembled code
    20  	}{
    21  		{
    22  			// string folding
    23  			`"a" + "b" + "c" + "d"`,
    24  			`constant "abcd"; return`,
    25  		},
    26  		{
    27  			// string folding with variable:
    28  			`"a" + "b" + x + "c" + "d"`,
    29  			`constant "ab"; predeclared x; plus; constant "cd"; plus; return`,
    30  		},
    31  		{
    32  			// list folding
    33  			`[1] + [2] + [3]`,
    34  			`constant 1; constant 2; constant 3; makelist<3>; return`,
    35  		},
    36  		{
    37  			// list folding with variable
    38  			`[1] + [2] + x + [3]`,
    39  			`constant 1; constant 2; makelist<2>; ` +
    40  				`predeclared x; plus; ` +
    41  				`constant 3; makelist<1>; plus; ` +
    42  				`return`,
    43  		},
    44  		{
    45  			// tuple folding
    46  			`() + (1,) + (2, 3)`,
    47  			`constant 1; constant 2; constant 3; maketuple<3>; return`,
    48  		},
    49  		{
    50  			// tuple folding with variable
    51  			`() + (1,) + x + (2, 3)`,
    52  			`constant 1; maketuple<1>; predeclared x; plus; ` +
    53  				`constant 2; constant 3; maketuple<2>; plus; ` +
    54  				`return`,
    55  		},
    56  	} {
    57  		expr, err := syntax.ParseExpr("in.star", test.src, 0)
    58  		if err != nil {
    59  			t.Errorf("#%d: %v", i, err)
    60  			continue
    61  		}
    62  		locals, err := resolve.Expr(expr, isPredeclared, isUniversal)
    63  		if err != nil {
    64  			t.Errorf("#%d: %v", i, err)
    65  			continue
    66  		}
    67  		got := disassemble(Expr(expr, "<expr>", locals).Toplevel)
    68  		if test.want != got {
    69  			t.Errorf("expression <<%s>> generated <<%s>>, want <<%s>>",
    70  				test.src, got, test.want)
    71  		}
    72  	}
    73  }
    74  
    75  // disassemble is a trivial disassembler tailored to the accumulator test.
    76  func disassemble(f *Funcode) string {
    77  	out := new(bytes.Buffer)
    78  	code := f.Code
    79  	for pc := 0; pc < len(code); {
    80  		op := Opcode(code[pc])
    81  		pc++
    82  		// TODO(adonovan): factor in common with interpreter.
    83  		var arg uint32
    84  		if op >= OpcodeArgMin {
    85  			for s := uint(0); ; s += 7 {
    86  				b := code[pc]
    87  				pc++
    88  				arg |= uint32(b&0x7f) << s
    89  				if b < 0x80 {
    90  					break
    91  				}
    92  			}
    93  		}
    94  
    95  		if out.Len() > 0 {
    96  			out.WriteString("; ")
    97  		}
    98  		fmt.Fprintf(out, "%s", op)
    99  		if op >= OpcodeArgMin {
   100  			switch op {
   101  			case CONSTANT:
   102  				switch x := f.Prog.Constants[arg].(type) {
   103  				case string:
   104  					fmt.Fprintf(out, " %q", x)
   105  				default:
   106  					fmt.Fprintf(out, " %v", x)
   107  				}
   108  			case LOCAL:
   109  				fmt.Fprintf(out, " %s", f.Locals[arg].Name)
   110  			case PREDECLARED:
   111  				fmt.Fprintf(out, " %s", f.Prog.Names[arg])
   112  			default:
   113  				fmt.Fprintf(out, "<%d>", arg)
   114  			}
   115  		}
   116  	}
   117  	return out.String()
   118  }