compiler: make stack depth calcuations work

Author: ncw

compile/compile.go

@@ -99,7 +99,6 @@ func Compile(str, filename, mode string, flags int, dont_inherit bool) py.Object
 	fmt.Println(ast.Dump(Ast))
 	code := &py.Code{
 		Filename:    filename,
-		Stacksize:   1,          // FIXME
 		Firstlineno: 1,          // FIXME
 		Name:        "<module>", // FIXME
 		Flags:       64,         // FIXME
@@ -121,6 +120,7 @@ func Compile(str, filename, mode string, flags int, dont_inherit bool) py.Object
 	}
 	c.Op(vm.RETURN_VALUE)
 	code.Code = c.OpCodes.Assemble()
+	code.Stacksize = int32(c.OpCodes.StackDepth())
 	return code
 }
 
@@ -496,16 +496,16 @@ func (is *Instructions) Add(i Instruction) {
 func (is Instructions) Pass(pass int) bool {
 	addr := uint32(0)
 	changed := false
-	for _, i := range is {
-		posChanged := i.SetPos(addr)
+	for i, instr := range is {
+		posChanged := instr.SetPos(i, addr)
 		changed = changed || posChanged
 		if pass > 0 {
 			// Only resolve addresses on 2nd pass
-			if resolver, ok := i.(Resolver); ok {
+			if resolver, ok := instr.(Resolver); ok {
 				resolver.Resolve()
 			}
 		}
-		addr += i.Size()
+		addr += instr.Size()
 	}
 	return changed
 }
@@ -527,31 +527,271 @@ done:
 	return string(out)
 }
 
+// Calculate number of arguments for CALL_FUNCTION etc
+func nArgs(o uint32) int {
+	return (int(o) & 0xFF) + 2*((int(o)>>8)&0xFF)
+}
+
+// Effect the opcode has on the stack
+func opcodeStackEffect(opcode byte, oparg uint32) int {
+	switch opcode {
+	case vm.POP_TOP:
+		return -1
+	case vm.ROT_TWO, vm.ROT_THREE:
+		return 0
+	case vm.DUP_TOP:
+		return 1
+	case vm.DUP_TOP_TWO:
+		return 2
+	case vm.UNARY_POSITIVE, vm.UNARY_NEGATIVE, vm.UNARY_NOT, vm.UNARY_INVERT:
+		return 0
+	case vm.SET_ADD, vm.LIST_APPEND:
+		return -1
+	case vm.MAP_ADD:
+		return -2
+	case vm.BINARY_POWER, vm.BINARY_MULTIPLY, vm.BINARY_MODULO, vm.BINARY_ADD, vm.BINARY_SUBTRACT, vm.BINARY_SUBSCR, vm.BINARY_FLOOR_DIVIDE, vm.BINARY_TRUE_DIVIDE:
+		return -1
+	case vm.INPLACE_FLOOR_DIVIDE, vm.INPLACE_TRUE_DIVIDE:
+		return -1
+	case vm.INPLACE_ADD, vm.INPLACE_SUBTRACT, vm.INPLACE_MULTIPLY, vm.INPLACE_MODULO:
+		return -1
+	case vm.STORE_SUBSCR:
+		return -3
+	case vm.STORE_MAP:
+		return -2
+	case vm.DELETE_SUBSCR:
+		return -2
+	case vm.BINARY_LSHIFT, vm.BINARY_RSHIFT, vm.BINARY_AND, vm.BINARY_XOR, vm.BINARY_OR:
+		return -1
+	case vm.INPLACE_POWER:
+		return -1
+	case vm.GET_ITER:
+		return 0
+	case vm.PRINT_EXPR:
+		return -1
+	case vm.LOAD_BUILD_CLASS:
+		return 1
+	case vm.INPLACE_LSHIFT, vm.INPLACE_RSHIFT, vm.INPLACE_AND, vm.INPLACE_XOR, vm.INPLACE_OR:
+		return -1
+	case vm.BREAK_LOOP:
+		return 0
+	case vm.SETUP_WITH:
+		return 7
+	case vm.WITH_CLEANUP:
+		return -1 /* XXX Sometimes more */
+	case vm.RETURN_VALUE:
+		return -1
+	case vm.IMPORT_STAR:
+		return -1
+	case vm.YIELD_VALUE:
+		return 0
+	case vm.YIELD_FROM:
+		return -1
+	case vm.POP_BLOCK:
+		return 0
+	case vm.POP_EXCEPT:
+		return 0 /* -3 except if bad bytecode */
+	case vm.END_FINALLY:
+		return -1 /* or -2 or -3 if exception occurred */
+	case vm.STORE_NAME:
+		return -1
+	case vm.DELETE_NAME:
+		return 0
+	case vm.UNPACK_SEQUENCE:
+		return int(oparg) - 1
+	case vm.UNPACK_EX:
+		return (int(oparg) & 0xFF) + (int(oparg) >> 8)
+	case vm.FOR_ITER:
+		return 1 /* or -1, at end of iterator */
+	case vm.STORE_ATTR:
+		return -2
+	case vm.DELETE_ATTR:
+		return -1
+	case vm.STORE_GLOBAL:
+		return -1
+	case vm.DELETE_GLOBAL:
+		return 0
+	case vm.LOAD_CONST:
+		return 1
+	case vm.LOAD_NAME:
+		return 1
+	case vm.BUILD_TUPLE, vm.BUILD_LIST, vm.BUILD_SET:
+		return 1 - int(oparg)
+	case vm.BUILD_MAP:
+		return 1
+	case vm.LOAD_ATTR:
+		return 0
+	case vm.COMPARE_OP:
+		return -1
+	case vm.IMPORT_NAME:
+		return -1
+	case vm.IMPORT_FROM:
+		return 1
+	case vm.JUMP_FORWARD, vm.JUMP_ABSOLUTE:
+		return 0
+	case vm.JUMP_IF_TRUE_OR_POP: /* -1 if jump not taken */
+		return 0
+	case vm.JUMP_IF_FALSE_OR_POP: /*  "" */
+		return 0
+	case vm.POP_JUMP_IF_FALSE, vm.POP_JUMP_IF_TRUE:
+		return -1
+	case vm.LOAD_GLOBAL:
+		return 1
+	case vm.CONTINUE_LOOP:
+		return 0
+	case vm.SETUP_LOOP:
+		return 0
+	case vm.SETUP_EXCEPT, vm.SETUP_FINALLY:
+		// can push 3 values for the new exception
+		// + 3 others for the previous exception state
+		return 6
+	case vm.LOAD_FAST:
+		return 1
+	case vm.STORE_FAST:
+		return -1
+	case vm.DELETE_FAST:
+		return 0
+
+	case vm.RAISE_VARARGS:
+		return -int(oparg)
+	case vm.CALL_FUNCTION:
+		return -nArgs(oparg)
+	case vm.CALL_FUNCTION_VAR, vm.CALL_FUNCTION_KW:
+		return -nArgs(oparg) - 1
+	case vm.CALL_FUNCTION_VAR_KW:
+		return -nArgs(oparg) - 2
+	case vm.MAKE_FUNCTION:
+		return -1 - nArgs(oparg) - ((int(oparg) >> 16) & 0xffff)
+	case vm.MAKE_CLOSURE:
+		return -2 - nArgs(oparg) - ((int(oparg) >> 16) & 0xffff)
+	case vm.BUILD_SLICE:
+		if oparg == 3 {
+			return -2
+		} else {
+			return -1
+		}
+	case vm.LOAD_CLOSURE:
+		return 1
+	case vm.LOAD_DEREF, vm.LOAD_CLASSDEREF:
+		return 1
+	case vm.STORE_DEREF:
+		return -1
+	case vm.DELETE_DEREF:
+		return 0
+	default:
+		panic("Unknown opcode in StackEffect")
+	}
+}
+
+// Recursive instruction walker to find max stack depth
+func (is Instructions) stackDepthWalk(baseIs Instructions, seen map[int]bool, startDepth map[int]int, depth int, maxdepth int) int {
+	// var i, target_depth, effect int
+	// var instr *struct instr
+	// if b.b_seen || b.b_startdepth >= depth {
+	// 	return maxdepth
+	// }
+	// b.b_seen = 1
+	// b.b_startdepth = depth
+	if len(is) == 0 {
+		return maxdepth
+	}
+	start := is[0].Number()
+	if seen[start] {
+		// We are processing this block already
+		return maxdepth
+	}
+	if d, ok := startDepth[start]; ok && d >= depth {
+		// We've processed this block with a larger depth already
+		return maxdepth
+	}
+	seen[start] = true
+	startDepth[start] = depth
+	for _, instr := range is {
+		depth += instr.StackEffect()
+		if depth > maxdepth {
+			maxdepth = depth
+		}
+		if depth < 0 {
+			panic("Stack depth negative")
+		}
+		jrel, isJrel := instr.(*JumpRel)
+		jabs, isJabs := instr.(*JumpAbs)
+		if isJrel || isJabs {
+			var oparg *OpArg
+			var dest *Label
+			if isJrel {
+				oparg = &jrel.OpArg
+				dest = jrel.Dest
+			} else {
+				oparg = &jabs.OpArg
+				dest = jabs.Dest
+			}
+			opcode := oparg.Op
+			target_depth := depth
+			if opcode == vm.FOR_ITER {
+				target_depth = depth - 2
+			} else if opcode == vm.SETUP_FINALLY || opcode == vm.SETUP_EXCEPT {
+				target_depth = depth + 3
+				if target_depth > maxdepth {
+					maxdepth = target_depth
+				}
+			} else if opcode == vm.JUMP_IF_TRUE_OR_POP || opcode == vm.JUMP_IF_FALSE_OR_POP {
+				depth = depth - 1
+			}
+			isTarget := baseIs[dest.Number():]
+			maxdepth = isTarget.stackDepthWalk(baseIs, seen, startDepth, target_depth, maxdepth)
+			if opcode == vm.JUMP_ABSOLUTE ||
+				opcode == vm.JUMP_FORWARD {
+				goto out // remaining code is dead
+			}
+		}
+	}
+out:
+	seen[start] = false
+	return maxdepth
+}
+
+// Find the flow path that needs the largest stack.  We assume that
+// cycles in the flow graph have no net effect on the stack depth.
+func (is Instructions) StackDepth() int {
+	return is.stackDepthWalk(is, make(map[int]bool), make(map[int]int), 0, 0)
+}
+
 type Instruction interface {
 	Pos() uint32
-	SetPos(uint32) bool
+	Number() int
+	SetPos(int, uint32) bool
 	Size() uint32
 	Output() []byte
+	StackEffect() int
 }
 
 type Resolver interface {
 	Resolve()
 }
 
 // Position
-type pos uint32
+type pos struct {
+	n uint32
+	p uint32
+}
+
+// Read instruction number
+func (p *pos) Number() int {
+	return int(p.n)
+}
 
 // Read position
 func (p *pos) Pos() uint32 {
-	return uint32(*p)
+	return p.p
 }
 
 // Set Position - returns changed
-func (p *pos) SetPos(newPos uint32) bool {
-	oldP := *p
-	newP := pos(newPos)
-	*p = newP
-	return oldP != newP
+func (p *pos) SetPos(number int, newPos uint32) bool {
+	p.n = uint32(number)
+	oldPos := p.p
+	p.p = newPos
+	return oldPos != newPos
 }
 
 // A plain opcode
@@ -570,6 +810,11 @@ func (o *Op) Output() []byte {
 	return []byte{byte(o.Op)}
 }
 
+// StackEffect
+func (o *Op) StackEffect() int {
+	return opcodeStackEffect(o.Op, 0)
+}
+
 // An opcode with argument
 type OpArg struct {
 	pos
@@ -595,6 +840,11 @@ func (o *OpArg) Output() []byte {
 	return out
 }
 
+// StackEffect
+func (o *OpArg) StackEffect() int {
+	return opcodeStackEffect(o.Op, o.Arg)
+}
+
 // A label
 type Label struct {
 	pos
@@ -610,6 +860,11 @@ func (o Label) Output() []byte {
 	return []byte{}
 }
 
+// StackEffect
+func (o *Label) StackEffect() int {
+	return 0
+}
+
 // An absolute JUMP with destination label
 type JumpAbs struct {
 	pos

compile/compile_test.go

@@ -87,6 +87,7 @@ func TestCompile(t *testing.T) {
 		if !ok {
 			t.Fatalf("Expecting *py.Code but got %T", codeObj)
 		}
+		t.Logf("Testing %q", test.in)
 		EqCode(t, &test.out, code)
 	}
 }