Add memory budget

Author: antonmedv

cmd/exe/main.go

@@ -14,6 +14,7 @@ import (
 	"github.com/antonmedv/expr/internal/conf"
 	"github.com/antonmedv/expr/optimizer"
 	"github.com/antonmedv/expr/parser"
+	"github.com/antonmedv/expr/vm"
 	"github.com/sanity-io/litter"
 )
 
@@ -118,7 +119,20 @@ func printDisassemble() {
 }
 
 func runProgram() {
-	out, err := expr.Eval(input(), nil)
+	tree, err := parser.Parse(input())
+	check(err)
+
+	_, err = checker.Check(tree, nil)
+	check(err)
+
+	if opt {
+		optimizer.Optimize(&tree.Node)
+	}
+
+	program, err := compiler.Compile(tree, nil)
+	check(err)
+
+	out, err := vm.Run(program, nil)
 	check(err)
 
 	litter.Dump(out)

optimizer/optimizer.go

@@ -237,6 +237,11 @@ func (*constRange) Exit(node *Node) {
 			if min, ok := n.Left.(*IntegerNode); ok {
 				if max, ok := n.Right.(*IntegerNode); ok {
 					size := max.Value - min.Value + 1
+					// In this case array is too big. Skip generation,
+					// and wait for memory budget detection on runtime.
+					if size > 1e6 {
+						return
+					}
 					value := make([]int, size)
 					for i := range value {
 						value[i] = min.Value + i

vm/runtime.go

@@ -218,15 +218,18 @@ func exponent(a, b interface{}) float64 {
 	return math.Pow(toFloat64(a), toFloat64(b))
 }
 
-func makeRange(a, b interface{}) []int {
+func makeRange(a, b interface{}, budget int) ([]int, int) {
 	min := toInt(a)
 	max := toInt(b)
 	size := max - min + 1
+	if size >= budget {
+		panic("memory budget exceeded")
+	}
 	rng := make([]int, size)
 	for i := range rng {
 		rng[i] = min + i
 	}
-	return rng
+	return rng, size
 }
 
 func toInt(a interface{}) int {

vm/vm.go

@@ -9,6 +9,10 @@ import (
 	"github.com/antonmedv/expr/file"
 )
 
+var (
+	MemoryBudget int = 1e6
+)
+
 func Run(program *Program, env interface{}) (out interface{}, err error) {
 	if program == nil {
 		return nil, fmt.Errorf("program is nil")
@@ -40,12 +44,15 @@ type VM struct {
 	debug     bool
 	step      chan struct{}
 	curr      chan int
+	memory    int
+	limit     int
 }
 
 func NewVM(debug bool) *VM {
 	vm := &VM{
 		stack: make([]interface{}, 0, 2),
 		debug: debug,
+		limit: MemoryBudget,
 	}
 	if vm.debug {
 		vm.step = make(chan struct{}, 0)
@@ -64,6 +71,10 @@ func (vm *VM) Run(program *Program, env interface{}) interface{} {
 			<-vm.step
 		}
 
+		if vm.memory >= vm.limit {
+			panic("memory budget exceeded")
+		}
+
 		vm.pp = vm.ip
 		vm.ip++
 		op := vm.bytecode[vm.pp]
@@ -198,7 +209,9 @@ func (vm *VM) Run(program *Program, env interface{}) interface{} {
 		case OpRange:
 			b := vm.pop()
 			a := vm.pop()
-			vm.push(makeRange(a, b))
+			c, size := makeRange(a, b, vm.limit)
+			vm.push(c)
+			vm.memory += size
 
 		case OpMatches:
 			b := vm.pop()
@@ -294,6 +307,7 @@ func (vm *VM) Run(program *Program, env interface{}) interface{} {
 				array[i] = vm.pop()
 			}
 			vm.push(array)
+			vm.memory += size
 
 		case OpMap:
 			size := vm.pop().(int)
@@ -304,6 +318,7 @@ func (vm *VM) Run(program *Program, env interface{}) interface{} {
 				m[key.(string)] = value
 			}
 			vm.push(m)
+			vm.memory += size
 
 		case OpLen:
 			vm.push(length(vm.current()))

vm/vm_test.go

@@ -94,3 +94,16 @@ func TestRun_helpers(t *testing.T) {
 		}
 	}
 }
+
+func TestRun_memory_budget(t *testing.T) {
+	input := `map(1..100, {map(1..100, {map(1..100, {0})})})`
+
+	tree, err := parser.Parse(input)
+	require.NoError(t, err)
+
+	program, err := compiler.Compile(tree, nil)
+	require.NoError(t, err)
+
+	_, err = vm.Run(program, nil)
+	require.Error(t, err)
+}