Add execution support for the tail-call proposal

README.md

@@ -68,6 +68,7 @@ and should work on the following platforms:
 |                      | [Sign-extension operators](https://github.com/WebAssembly/spec/blob/master/proposals/sign-extension-ops/Overview.md) | ✅ Implemented |
 |                      | [Non-trapping float-to-int conversions](https://github.com/WebAssembly/nontrapping-float-to-int-conversions/blob/main/proposals/nontrapping-float-to-int-conversion/Overview.md) | ✅ Implemented |
 |                      | [Memory64](https://github.com/WebAssembly/memory64/blob/main/proposals/memory64/Overview.md) | ✅ Implemented |
+|                      | [Tail call](https://github.com/WebAssembly/tail-call/blob/master/proposals/tail-call/Overview.md) | ✅ Implemented |
 |                      | [Threads and atomics](https://github.com/WebAssembly/threads/blob/master/proposals/threads/Overview.md) | 🚧 Parser implemented |
 | WASI | WASI Preview 1 | ✅ Implemented |
 

Sources/WasmKit/Execution/Execution.swift

@@ -78,27 +78,32 @@ struct Execution {
         return Backtrace(symbols: symbols)
     }
 
+    private func initializeConstSlots(
+        sp: Sp, iseq: InstructionSequence,
+        numberOfNonParameterLocals: Int
+    ) {
+        // Initialize the locals with zeros (all types of value have the same representation)
+        sp.initialize(repeating: UntypedValue.default.storage, count: numberOfNonParameterLocals)
+        if let constants = iseq.constants.baseAddress {
+            let count = iseq.constants.count
+            sp.advanced(by: numberOfNonParameterLocals).withMemoryRebound(to: UntypedValue.self, capacity: count) {
+                $0.initialize(from: constants, count: count)
+            }
+        }
+    }
+
     /// Pushes a new call frame to the VM stack.
     @inline(__always)
-    mutating func pushFrame(
+    func pushFrame(
         iseq: InstructionSequence,
         function: EntityHandle<WasmFunctionEntity>,
         numberOfNonParameterLocals: Int,
         sp: Sp, returnPC: Pc,
         spAddend: VReg
     ) throws -> Sp {
         let newSp = sp.advanced(by: Int(spAddend))
-        guard newSp.advanced(by: iseq.maxStackHeight) < stackEnd else {
-            throw Trap(.callStackExhausted)
-        }
-        // Initialize the locals with zeros (all types of value have the same representation)
-        newSp.initialize(repeating: UntypedValue.default.storage, count: numberOfNonParameterLocals)
-        if let constants = iseq.constants.baseAddress {
-            let count = iseq.constants.count
-            newSp.advanced(by: numberOfNonParameterLocals).withMemoryRebound(to: UntypedValue.self, capacity: count) {
-                $0.initialize(from: constants, count: count)
-            }
-        }
+        try checkStackBoundary(newSp.advanced(by: iseq.maxStackHeight))
+        initializeConstSlots(sp: newSp, iseq: iseq, numberOfNonParameterLocals: numberOfNonParameterLocals)
         newSp.previousSP = sp
         newSp.returnPC = returnPC
         newSp.currentFunction = function
@@ -107,7 +112,7 @@ struct Execution {
 
     /// Pops the current frame from the VM stack.
     @inline(__always)
-    mutating func popFrame(sp: inout Sp, pc: inout Pc, md: inout Md, ms: inout Ms) {
+    func popFrame(sp: inout Sp, pc: inout Pc, md: inout Md, ms: inout Ms) {
         let oldSp = sp
         sp = oldSp.previousSP.unsafelyUnwrapped
         pc = oldSp.returnPC.unsafelyUnwrapped
@@ -505,48 +510,114 @@ extension Execution {
         self.trap = (rawError, sp)
     }
 
+    @inline(__always)
+    func checkStackBoundary(_ sp: Sp) throws {
+        guard sp < stackEnd else { throw Trap(.callStackExhausted) }
+    }
+
     /// Returns the new program counter and stack pointer.
     @inline(never)
-    mutating func invoke(
+    func invoke(
         function: InternalFunction,
         callerInstance: InternalInstance?,
         spAddend: VReg,
         sp: Sp, pc: Pc, md: inout Md, ms: inout Ms
     ) throws -> (Pc, Sp) {
         if function.isWasm {
-            let function = function.wasm
-            let iseq = try function.ensureCompiled(store: store)
-
-            let newSp = try pushFrame(
-                iseq: iseq,
-                function: function,
-                numberOfNonParameterLocals: function.numberOfNonParameterLocals,
-                sp: sp,
-                returnPC: pc,
-                spAddend: spAddend
+            return try invokeWasmFunction(
+                function: function.wasm, callerInstance: callerInstance,
+                spAddend: spAddend, sp: sp, pc: pc, md: &md, ms: &ms
             )
-            Execution.CurrentMemory.mayUpdateCurrentInstance(
-                instance: function.instance,
-                from: callerInstance, md: &md, ms: &ms
-            )
-            return (iseq.baseAddress, newSp)
         } else {
-            let function = function.host
-            let resolvedType = store.value.engine.resolveType(function.type)
-            let layout = FrameHeaderLayout(type: resolvedType)
-            let parameters = resolvedType.parameters.enumerated().map { (i, type) in
-                sp[spAddend + layout.paramReg(i)].cast(to: type)
-            }
-            let instance = self.currentInstance(sp: sp)
-            let caller = Caller(
-                instanceHandle: instance,
-                store: store.value
+            try invokeHostFunction(function: function.host, sp: sp, spAddend: spAddend)
+            return (pc, sp)
+        }
+    }
+
+    @inline(never)
+    func tailInvoke(
+        function: InternalFunction,
+        callerInstance: InternalInstance?,
+        sp: Sp, pc: Pc, md: inout Md, ms: inout Ms
+    ) throws -> (Pc, Sp) {
+        if function.isWasm {
+            return try tailInvokeWasmFunction(
+                function: function.wasm, callerInstance: callerInstance,
+                sp: sp, md: &md, ms: &ms
             )
-            let results = try function.implementation(caller, Array(parameters))
-            for (index, result) in results.enumerated() {
-                sp[spAddend + layout.returnReg(index)] = UntypedValue(result)
-            }
+        } else {
+            try invokeHostFunction(function: function.host, sp: sp, spAddend: 0)
             return (pc, sp)
         }
     }
+
+    /// Executes the given wasm function while overwriting the current frame.
+    ///
+    /// Precondition: The frame header must be already resized to be compatible
+    /// with the callee's frame header layout.
+    @inline(__always)
+    private func tailInvokeWasmFunction(
+        function: EntityHandle<WasmFunctionEntity>,
+        callerInstance: InternalInstance?,
+        sp: Sp, md: inout Md, ms: inout Ms
+    ) throws -> (Pc, Sp) {
+        let iseq = try function.ensureCompiled(store: store)
+        try checkStackBoundary(sp.advanced(by: iseq.maxStackHeight))
+        sp.currentFunction = function
+
+        initializeConstSlots(sp: sp, iseq: iseq, numberOfNonParameterLocals: function.numberOfNonParameterLocals)
+
+        Execution.CurrentMemory.mayUpdateCurrentInstance(
+            instance: function.instance,
+            from: callerInstance, md: &md, ms: &ms
+        )
+        return (iseq.baseAddress, sp)
+    }
+
+    /// Executes the given WebAssembly function.
+    @inline(__always)
+    private func invokeWasmFunction(
+        function: EntityHandle<WasmFunctionEntity>,
+        callerInstance: InternalInstance?,
+        spAddend: VReg,
+        sp: Sp, pc: Pc, md: inout Md, ms: inout Ms
+    ) throws -> (Pc, Sp) {
+        let iseq = try function.ensureCompiled(store: store)
+
+        let newSp = try pushFrame(
+            iseq: iseq,
+            function: function,
+            numberOfNonParameterLocals: function.numberOfNonParameterLocals,
+            sp: sp,
+            returnPC: pc,
+            spAddend: spAddend
+        )
+        Execution.CurrentMemory.mayUpdateCurrentInstance(
+            instance: function.instance,
+            from: callerInstance, md: &md, ms: &ms
+        )
+        return (iseq.baseAddress, newSp)
+    }
+
+    /// Executes the given host function.
+    ///
+    /// Note that this function does not modify neither the positions of the
+    /// stack pointer nor the program counter.
+    @inline(never)
+    private func invokeHostFunction(function: EntityHandle<HostFunctionEntity>, sp: Sp, spAddend: VReg) throws {
+        let resolvedType = store.value.engine.resolveType(function.type)
+        let layout = FrameHeaderLayout(type: resolvedType)
+        let parameters = resolvedType.parameters.enumerated().map { (i, type) in
+            sp[spAddend + layout.paramReg(i)].cast(to: type)
+        }
+        let instance = self.currentInstance(sp: sp)
+        let caller = Caller(
+            instanceHandle: instance,
+            store: store.value
+        )
+        let results = try function.implementation(caller, Array(parameters))
+        for (index, result) in results.enumerated() {
+            sp[spAddend + layout.returnReg(index)] = UntypedValue(result)
+        }
+    }
 }

Sources/WasmKit/Execution/Instructions/Control.swift

@@ -112,11 +112,11 @@ extension Execution {
     @inline(never)
     private func prepareForIndirectCall(
         sp: Sp, tableIndex: TableIndex, expectedType: InternedFuncType,
-        callIndirectOperand: Instruction.CallIndirectOperand
+        address: VReg
     ) throws -> (InternalFunction, InternalInstance) {
         let callerInstance = currentInstance(sp: sp)
         let table = callerInstance.tables[Int(tableIndex)]
-        let value = sp[callIndirectOperand.index].asAddressOffset(table.limits.isMemory64)
+        let value = sp[address].asAddressOffset(table.limits.isMemory64)
         let elementIndex = Int(value)
         guard elementIndex < table.elements.count else {
             throw Trap(.tableOutOfBounds(elementIndex))
@@ -141,7 +141,7 @@ extension Execution {
         var pc = pc
         let (function, callerInstance) = try prepareForIndirectCall(
             sp: sp, tableIndex: immediate.tableIndex, expectedType: immediate.type,
-            callIndirectOperand: immediate
+            address: immediate.index
         )
         (pc, sp) = try invoke(
             function: function,
@@ -152,6 +152,65 @@ extension Execution {
         return pc.next()
     }
 
+    mutating func returnCall(sp: inout Sp, pc: Pc, md: inout Md, ms: inout Ms, immediate: Instruction.ReturnCallOperand) throws -> (Pc, CodeSlot) {
+        var pc = pc
+        (pc, sp) = try tailInvoke(
+            function: immediate.callee,
+            callerInstance: currentInstance(sp: sp),
+            sp: sp, pc: pc, md: &md, ms: &ms
+        )
+        return pc.next()
+    }
+
+    mutating func returnCallIndirect(sp: inout Sp, pc: Pc, md: inout Md, ms: inout Ms, immediate: Instruction.ReturnCallIndirectOperand) throws -> (Pc, CodeSlot) {
+        var pc = pc
+        let (function, callerInstance) = try prepareForIndirectCall(
+            sp: sp, tableIndex: immediate.tableIndex, expectedType: immediate.type,
+            address: immediate.index
+        )
+        (pc, sp) = try tailInvoke(
+            function: function,
+            callerInstance: callerInstance,
+            sp: sp, pc: pc, md: &md, ms: &ms
+        )
+        return pc.next()
+    }
+
+    mutating func resizeFrameHeader(sp: inout Sp, immediate: Instruction.ResizeFrameHeaderOperand) throws {
+        // The params/results space are resized by `delta` slots and the rest of the
+        // frame is copied to the new location. See the following diagram for the
+        // layout of the frame before and after the resize operation:
+        //
+        //
+        //              |--------BEFORE-------|   |--------AFTER--------|
+        //              |  Params  | Results  |   |  Params  | Results  |
+        //              |  ...     |   ...    |   |  ...     |   ...    |
+        // Old Header ->|---------------------|\  |  ...     |   ...    |              -+
+        //              |         Sp          | \ |  ...     |   ...    |               | delta
+        //              |---------------------|  \|---------------------|<- New Header -+  -+
+        //              |         Pc          |   |         Sp          |                   |
+        //              |---------------------|   |---------------------|                   |
+        //              |     Current Func    | C |         Pc          |                   |
+        //     Old Sp ->|---------------------| O |---------------------|                   |
+        //              |       Locals        | P |     Current Func    |                   |
+        //              |        ...          | Y |---------------------|<- New Sp          |
+        //              |---------------------|   |       Locals        |                   | sizeToCopy
+        //              |        Consts       |   |        ...          |                   |
+        //              |        ...          |   |---------------------|                   |
+        //              |---------------------|   |        Consts       |                   |
+        //              |     Value Stack     |   |        ...          |                   |
+        //              |        ...          |   |---------------------|                   |
+        //              |---------------------|\  |     Value Stack     |                   |
+        //                                      \ |        ...          |                   |
+        //                                       \|---------------------|                  -+
+        let newSp = sp.advanced(by: Int(immediate.delta))
+        try checkStackBoundary(newSp)
+        let oldFrameHeader = sp.advanced(by: -FrameHeaderLayout.numberOfSavingSlots)
+        let newFrameHeader = newSp.advanced(by: -FrameHeaderLayout.numberOfSavingSlots)
+        newFrameHeader.update(from: oldFrameHeader, count: Int(immediate.sizeToCopy))
+        sp = newSp
+    }
+
     mutating func onEnter(sp: Sp, immediate: Instruction.OnEnterOperand) {
         let function = currentInstance(sp: sp).functions[Int(immediate)]
         self.store.value.engine.interceptor?.onEnterFunction(