optimize vm allocation of function arguments

Author: diegommm

vm/vm.go

@@ -83,6 +83,8 @@ func (vm *VM) Run(program *Program, env any) (_ any, err error) {
 	vm.memory = 0
 	vm.ip = 0
 
+	var fnArgsBuf []any
+
 	for vm.ip < len(program.Bytecode) {
 		if debug && vm.debug {
 			<-vm.step
@@ -355,62 +357,47 @@ func (vm *VM) Run(program *Program, env any) (_ any, err error) {
 			vm.push(out)
 
 		case OpCall1:
-			a := vm.pop()
-			out, err := program.functions[arg](a)
+			args := vm.getArgsForFunc(&fnArgsBuf, program, 1)
+			out, err := program.functions[arg](args...)
 			if err != nil {
 				panic(err)
 			}
 			vm.push(out)
 
 		case OpCall2:
-			b := vm.pop()
-			a := vm.pop()
-			out, err := program.functions[arg](a, b)
+			args := vm.getArgsForFunc(&fnArgsBuf, program, 2)
+			out, err := program.functions[arg](args...)
 			if err != nil {
 				panic(err)
 			}
 			vm.push(out)
 
 		case OpCall3:
-			c := vm.pop()
-			b := vm.pop()
-			a := vm.pop()
-			out, err := program.functions[arg](a, b, c)
+			args := vm.getArgsForFunc(&fnArgsBuf, program, 3)
+			out, err := program.functions[arg](args...)
 			if err != nil {
 				panic(err)
 			}
 			vm.push(out)
 
 		case OpCallN:
 			fn := vm.pop().(Function)
-			size := arg
-			in := make([]any, size)
-			for i := int(size) - 1; i >= 0; i-- {
-				in[i] = vm.pop()
-			}
-			out, err := fn(in...)
+			args := vm.getArgsForFunc(&fnArgsBuf, program, arg)
+			out, err := fn(args...)
 			if err != nil {
 				panic(err)
 			}
 			vm.push(out)
 
 		case OpCallFast:
 			fn := vm.pop().(func(...any) any)
-			size := arg
-			in := make([]any, size)
-			for i := int(size) - 1; i >= 0; i-- {
-				in[i] = vm.pop()
-			}
-			vm.push(fn(in...))
+			args := vm.getArgsForFunc(&fnArgsBuf, program, arg)
+			vm.push(fn(args...))
 
 		case OpCallSafe:
 			fn := vm.pop().(SafeFunction)
-			size := arg
-			in := make([]any, size)
-			for i := int(size) - 1; i >= 0; i-- {
-				in[i] = vm.pop()
-			}
-			out, mem, err := fn(in...)
+			args := vm.getArgsForFunc(&fnArgsBuf, program, arg)
+			out, mem, err := fn(args...)
 			if err != nil {
 				panic(err)
 			}
@@ -609,6 +596,56 @@ func (vm *VM) scope() *Scope {
 	return vm.Scopes[len(vm.Scopes)-1]
 }
 
+// getArgsForFunc lazily initializes the buffer the first time it is called for
+// a given program (thus, it also needs "program" to run). It will
+// take "needed" elements from the buffer and populate them with vm.pop() in
+// reverse order. Because the estimation can fall short, this function can
+// occasionally make a new allocation.
+func (vm *VM) getArgsForFunc(bufPtr *[]any, program *Program, needed int) []any {
+	// Step 1: fix estimations and preallocate
+	if *bufPtr == nil {
+		estimatedFnArgsCount := estimateFnArgsCount(program)
+		if estimatedFnArgsCount < needed {
+			// in the case that the first call is for example OpCallN with a large
+			// number of arguments, then make sure we will be able to serve them at
+			// least.
+			estimatedFnArgsCount = needed
+		}
+
+		// in the case that we are preparing the arguments for the first
+		// function call of the program, then *bufPtr will be nil, so we
+		// initialize it. We delay this initial allocation here because a
+		// program could have many function calls but exit earlier than the
+		// first call, so in that case we avoid allocating unnecessarily
+		*bufPtr = make([]any, estimatedFnArgsCount)
+	}
+
+	// Step 2: get the final slice that will be returned
+	var buf []any
+	if len(*bufPtr) >= needed {
+		// in this case, we are successfully using the single preallocation. We
+		// use the full slice expression [low : high : max] because in that way
+		// a function that receives this slice as variadic arguments will not be
+		// able to make modifications to contiguous elements with append(). If
+		// they call append on their variadic arguments they will make a new
+		// allocation.
+		buf = (*bufPtr)[:needed:needed]
+		*bufPtr = (*bufPtr)[needed:] // advance the buffer
+	} else {
+		// if we have been making calls to something like OpCallN with many more
+		// arguments than what we estimated, then we will need to allocate
+		// separately
+		buf = make([]any, needed)
+	}
+
+	// Step 3: populate the final slice bulk copying from the stack. This is the
+	// exact order and copy() is a highly optimized operation
+	copy(buf, vm.Stack[len(vm.Stack)-needed:])
+	vm.Stack = vm.Stack[:len(vm.Stack)-needed]
+
+	return buf
+}
+
 func (vm *VM) Step() {
 	vm.step <- struct{}{}
 }
@@ -623,3 +660,31 @@ func clearSlice[S ~[]E, E any](s S) {
 		s[i] = zero // clear mem, optimized by the compiler, in Go 1.21 the "clear" builtin can be used
 	}
 }
+
+// estimateFnArgsCount inspects a *Program and estimates how many function
+// arguments will be required to run it.
+func estimateFnArgsCount(program *Program) int {
+	// Implementation note: a program will not necessarily go through all
+	// operations, but this is just an estimation
+	var count int
+	for _, op := range program.Bytecode {
+		switch op {
+		case OpCall1:
+			count++
+		case OpCall2:
+			count += 2
+		case OpCall3:
+			count += 3
+		case OpCallN:
+			// we don't know exactly but we know at least 4, so be conservative
+			// as this is only an optimization and we also want to avoid
+			// excessive preallocation
+			count += 4
+		case OpCallFast, OpCallSafe:
+			// here we don't know either, but we can guess it could be common to
+			// receive up to 3 arguments in a function
+			count += 3
+		}
+	}
+	return count
+}