Add a MaxVariables limits

engine/builtin_test.go

@@ -952,7 +952,7 @@ func TestTermVariables(t *testing.T) {
 func TestOp(t *testing.T) {
 	t.Run("insert", func(t *testing.T) {
 		t.Run("atom", func(t *testing.T) {
-			varCounter = 1
+			varCounter.count = 1
 
 			vm := VM{_operators: newOperators()}
 			vm.getOperators().define(900, operatorSpecifierXFX, NewAtom(`+++`))

engine/text_test.go

@@ -24,7 +24,7 @@ func mustOpen(fs fs.FS, name string) fs.File {
 }
 
 func TestVM_Compile(t *testing.T) {
-	varCounter = 1
+	varCounter.count = 1
 
 	tests := []struct {
 		title  string
@@ -483,7 +483,7 @@ bar(b).
 	for _, tt := range tests {
 		t.Run(tt.title, func(t *testing.T) {
 			var vm VM
-			varCounter = 1 // Global var cause issues in testing environment that call in randomly order for checking equality between procedure clause args
+			varCounter.count = 1 // Global var cause issues in testing environment that call in randomly order for checking equality between procedure clause args
 
 			vm.getOperators().define(1200, operatorSpecifierXFX, atomIf)
 			vm.getOperators().define(1200, operatorSpecifierXFX, atomArrow)

engine/variable.go

@@ -1,24 +1,40 @@
 package engine
 
 import (
+	"errors"
 	"fmt"
 	"io"
-	"sync/atomic"
+	"sync"
 )
 
-var varCounter int64
+var errMaxVariables = errors.New("maximum number of variables reached")
+
+var maxVariables uint64
+var varCounter = struct {
+	sync.Mutex
+	count uint64
+}{
+	count: 0,
+}
 
 func lastVariable() Variable {
-	return Variable(varCounter)
+	defer varCounter.Unlock()
+	varCounter.Lock()
+	return Variable(varCounter.count)
 }
 
 // Variable is a prolog variable.
 type Variable int64
 
 // NewVariable creates a new anonymous variable.
 func NewVariable() Variable {
-	n := atomic.AddInt64(&varCounter, 1)
-	return Variable(n)
+	defer varCounter.Unlock()
+	varCounter.Lock()
+	if maxVariables != 0 && varCounter.count >= maxVariables {
+		panic(errMaxVariables)
+	}
+	varCounter.count++
+	return Variable(varCounter.count)
 }
 
 func (v Variable) WriteTerm(w io.Writer, opts *WriteOptions, env *Env) error {

engine/variable_test.go

@@ -141,3 +141,39 @@ func Test_freeVariablesSet(t *testing.T) {
 		assert.Equal(t, tt.fv, newFreeVariablesSet(tt.t, tt.v, nil))
 	}
 }
+
+func Test_maxVariables(t *testing.T) {
+	tests := []struct {
+		title         string
+		init          func()
+		max           uint64
+		expectedCount uint64
+		shouldPanic   bool
+	}{
+		{title: "no limits", init: func() {
+			NewVariable()
+			NewVariable()
+		}, max: 0, expectedCount: 2},
+		{title: "limit", init: func() {
+			NewVariable()
+			NewVariable()
+		}, max: 2, expectedCount: 2},
+		{title: "limit reached", init: func() {
+			NewVariable()
+			NewVariable()
+		}, max: 1, expectedCount: 1, shouldPanic: true},
+	}
+
+	for _, tt := range tests {
+		varCounter.count = 0 // reset at each test
+
+		maxVariables = tt.max
+
+		if tt.shouldPanic {
+			assert.Panics(t, tt.init)
+		} else {
+			tt.init()
+		}
+		assert.Equal(t, varCounter.count, tt.expectedCount)
+	}
+}

engine/vm.go

@@ -70,6 +70,9 @@ type VM struct {
 	streams       streams
 	input, output *Stream
 
+	// Limits
+	maxVariables uint64
+
 	// Misc
 	debug bool
 }
@@ -276,9 +279,16 @@ func (vm *VM) SetUserOutput(s *Stream) {
 	vm.output = s
 }
 
+// SetMaxVariables sets the maximum number of variables that the VM can create.
+// Zero value mean no limits
+func (vm *VM) SetMaxVariables(n uint64) {
+	vm.maxVariables = n
+	maxVariables = n
+}
+
 // ResetEnv is used to reset all global variable
 func (vm *VM) ResetEnv() {
-	varCounter = 0
+	varCounter.count = 0
 	varContext = NewVariable()
 	rootContext = NewAtom("root")
 	rootEnv = &Env{
@@ -287,6 +297,7 @@ func (vm *VM) ResetEnv() {
 			value: rootContext,
 		},
 	}
+	maxVariables = vm.maxVariables
 }
 
 func (vm *VM) getProcedure(p procedureIndicator) (procedure, bool) {

engine/vm_test.go

@@ -270,6 +270,15 @@ func TestVM_SetUserOutput(t *testing.T) {
 	})
 }
 
+func TestVM_SetMaxVariables(t *testing.T) {
+	t.Run("ok", func(t *testing.T) {
+		var vm VM
+		vm.SetMaxVariables(10)
+		assert.Equal(t, uint64(10), maxVariables)
+		assert.Equal(t, uint64(10), vm.maxVariables)
+	})
+}
+
 func TestProcedureIndicator_Apply(t *testing.T) {
 	t.Run("ok", func(t *testing.T) {
 		c, err := procedureIndicator{name: NewAtom("foo"), arity: 2}.Apply(NewAtom("a"), NewAtom("b"))
@@ -289,7 +298,9 @@ func TestProcedureIndicator_Apply(t *testing.T) {
 
 func TestVM_ResetEnv(t *testing.T) {
 	var vm VM
-	varCounter = 10
+	vm.SetMaxVariables(20)
+
+	varCounter.count = 10
 	varContext = NewVariable()
 	rootContext = NewAtom("non-root")
 	rootEnv = &Env{
@@ -298,13 +309,15 @@ func TestVM_ResetEnv(t *testing.T) {
 			value: NewAtom("non-root"),
 		},
 	}
+	maxVariables = 30
 
 	t.Run("Reset environment", func(t *testing.T) {
 		vm.ResetEnv()
 
-		assert.Equal(t, int64(1), varCounter) // 1 because NewVariable() is called in ResetEnv()
+		assert.Equal(t, uint64(1), varCounter.count) // 1 because NewVariable() is called in ResetEnv()
 		assert.Equal(t, "root", rootContext.String())
 		assert.Equal(t, newEnvKey(varContext), rootEnv.binding.key)
 		assert.Equal(t, NewAtom("root"), rootEnv.binding.value)
+		assert.Equal(t, uint64(20), maxVariables)
 	})
 }