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