@@ -10,43 +10,41 @@ import (
 ".com/expr-lang/expr/vm/runtime"
 )
 
-const (
-// DefaultMemoryBudget represents an upper limit of memory usage
+var (
+// DefaultMemoryBudget represents default maximum allowed memory usage by the vm.VM.
 DefaultMemoryBudget uint = 1e6
 
-// DefaultMaxNodes represents an upper limit of AST nodes
-DefaultMaxNodes uint = 10000
+// DefaultMaxNodes represents default maximum allowed AST nodes by the compiler.
+DefaultMaxNodes uint = 1e4
 )
 
 type FunctionsTable map[string]*builtin.Function
 
 type Config struct {
-EnvObject any
-Env nature.Nature
-Expect reflect.Kind
-ExpectAny bool
-Optimize bool
-Strict bool
-Profile bool
-MaxNodes uint
-MemoryBudget uint
-ConstFns map[string]reflect.Value
-Visitors []ast.Visitor
-Functions FunctionsTable
-Builtins FunctionsTable
-Disabled map[string]bool // disabled builtins
+EnvObject any
+Env nature.Nature
+Expect reflect.Kind
+ExpectAny bool
+Optimize bool
+Strict bool
+Profile bool
+MaxNodes uint
+ConstFns map[string]reflect.Value
+Visitors []ast.Visitor
+Functions FunctionsTable
+Builtins FunctionsTable
+Disabled map[string]bool // disabled builtins
 }
 
 // CreateNew creates new config with default values.
 func CreateNew() *Config {
 c := &Config{
-Optimize: true,
-MaxNodes: DefaultMaxNodes,
-MemoryBudget: DefaultMemoryBudget,
-ConstFns: make(map[string]reflect.Value),
-Functions: make(map[string]*builtin.Function),
-Builtins: make(map[string]*builtin.Function),
-Disabled: make(map[string]bool),
+Optimize: true,
+MaxNodes: DefaultMaxNodes,
+ConstFns: make(map[string]reflect.Value),
+Functions: make(map[string]*builtin.Function),
+Builtins: make(map[string]*builtin.Function),
+Disabled: make(map[string]bool),
 }
 for _, f := range builtin.Builtins {
 c.Builtins[f.Name] = f

@@ -195,6 +195,15 @@ func Timezone(name string) Option {
 })
 }
 
+// MaxNodes sets the maximum number of nodes allowed in the expression.
+// By default, the maximum number of nodes is conf.DefaultMaxNodes.
+// If MaxNodes is set to 0, the node budget check is disabled.
+func MaxNodes(n uint) Option {
+return func(c *conf.Config) {
+c.MaxNodes = n
+}
+}
+
 // Compile parses and compiles given input expression to bytecode program.
 func Compile(input string, ops ...Option) (*vm.Program, error) {
 config := conf.CreateNew()

@@ -10,9 +10,11 @@ import (
 "testing"
 "time"
 
+".com/expr-lang/expr/conf"
 ".com/expr-lang/expr/internal/testify/assert"
 ".com/expr-lang/expr/internal/testify/require"
 ".com/expr-lang/expr/types"
+".com/expr-lang/expr/vm"
 
 ".com/expr-lang/expr"
 ".com/expr-lang/expr/ast"
@@ -2225,26 +2227,6 @@ func TestEval_slices_out_of_bound(t *testing.T) {
 }
 }
 
-func TestMemoryBudget(t *testing.T) {
-tests := []struct {
-code string
-}{
-{`map(1..100, {map(1..100, {map(1..100, {0})})})`},
-{`len(1..10000000)`},
-}
-
-for _, tt := range tests {
-t.Run(tt.code, func(t *testing.T) {
-program, err := expr.Compile(tt.code)
-require.NoError(t, err, "compile error")
-
-_, err = expr.Run(program, nil)
-assert.Error(t, err, "run error")
-assert.Contains(t, err.Error(), "memory budget exceeded")
-})
-}
-}
-
 func TestExpr_custom_tests(t *testing.T) {
 f, err := os.Open("custom_tests.json")
 if os.IsNotExist(err) {
@@ -2731,3 +2713,55 @@ func TestIssue785_get_nil(t *testing.T) {
 })
 }
 }
+
+func TestMaxNodes(t *testing.T) {
+maxNodes := uint(100)
+
+code := ""
+for i := 0; i < int(maxNodes); i++ {
+code += "1; "
+}
+
+_, err := expr.Compile(code, expr.MaxNodes(maxNodes))
+require.Error(t, err)
+assert.Contains(t, err.Error(), "exceeds maximum allowed nodes")
+
+_, err = expr.Compile(code, expr.MaxNodes(maxNodes+1))
+require.NoError(t, err)
+}
+
+func TestMaxNodesDisabled(t *testing.T) {
+code := ""
+for i := 0; i < 2*int(conf.DefaultMaxNodes); i++ {
+code += "1; "
+}
+
+_, err := expr.Compile(code, expr.MaxNodes(0))
+require.NoError(t, err)
+}
+
+func TestMemoryBudget(t *testing.T) {
+tests := []struct {
+code string
+max int
+}{
+{`map(1..100, {map(1..100, {map(1..100, {0})})})`, -1},
+{`len(1..10000000)`, -1},
+{`1..100`, 100},
+}
+
+for _, tt := range tests {
+t.Run(tt.code, func(t *testing.T) {
+program, err := expr.Compile(tt.code)
+require.NoError(t, err, "compile error")
+
+vm := vm.VM{}
+if tt.max > 0 {
+vm.MemoryBudget = uint(tt.max)
+}
+_, err = vm.Run(program, nil)
+require.Error(t, err, "run error")
+assert.Contains(t, err.Error(), "memory budget exceeded")
+})
+}
+}

@@ -75,7 +75,6 @@ func (vm *VM) Run(program *Program, env any) (_ any, err error) {
 if len(vm.Variables) < program.variables {
 vm.Variables = make([]any, program.variables)
 }
-
 if vm.MemoryBudget == 0 {
 vm.MemoryBudget = conf.DefaultMemoryBudget
 }