Unwinding Follow-up. Fix a bug in JitUnwindWindows where ... (#1238)

... in case of "Vector" unwind codes the remaining unwind codes could be corrupted.
Nits.

ARMeilleure/CodeGen/Unwinding/UnwindInfo.cs

@@ -3,16 +3,12 @@ namespace ARMeilleure.CodeGen.Unwinding
     struct UnwindInfo
     {
         public UnwindPushEntry[] PushEntries { get; }
+        public int PrologSize { get; }
 
-        public int PrologueSize { get; }
-
-        public int FixedAllocSize { get; }
-
-        public UnwindInfo(UnwindPushEntry[] pushEntries, int prologueSize, int fixedAllocSize)
+        public UnwindInfo(UnwindPushEntry[] pushEntries, int prologSize)
         {
-            PushEntries    = pushEntries;
-            PrologueSize   = prologueSize;
-            FixedAllocSize = fixedAllocSize;
+            PushEntries = pushEntries;
+            PrologSize = prologSize;
         }
     }
 }

ARMeilleure/CodeGen/Unwinding/UnwindPseudoOperation.cs

@@ -0,0 +1,11 @@
+namespace ARMeilleure.CodeGen.Unwinding
+{
+    enum UnwindPseudoOp
+    {
+        PushReg,
+        SetFrame,
+        AllocStack,
+        SaveReg,
+        SaveXmm128
+    }
+}

ARMeilleure/CodeGen/Unwinding/UnwindPushEntry.cs

@@ -1,20 +1,18 @@
-using ARMeilleure.IntermediateRepresentation;
-
 namespace ARMeilleure.CodeGen.Unwinding
 {
     struct UnwindPushEntry
     {
-        public int Index { get; }
-
-        public RegisterType Type { get; }
-
-        public int StreamEndOffset { get; }
+        public UnwindPseudoOp PseudoOp { get; }
+        public int PrologOffset { get; }
+        public int RegIndex { get; }
+        public int StackOffsetOrAllocSize { get; }
 
-        public UnwindPushEntry(int index, RegisterType type, int streamEndOffset)
+        public UnwindPushEntry(UnwindPseudoOp pseudoOp, int prologOffset, int regIndex = -1, int stackOffsetOrAllocSize = -1)
         {
-            Index           = index;
-            Type            = type;
-            StreamEndOffset = streamEndOffset;
+            PseudoOp = pseudoOp;
+            PrologOffset = prologOffset;
+            RegIndex = regIndex;
+            StackOffsetOrAllocSize = stackOffsetOrAllocSize;
         }
     }
 }

ARMeilleure/CodeGen/X86/CodeGenerator.cs

@@ -1525,30 +1525,27 @@ namespace ARMeilleure.CodeGen.X86
             context.Assembler.Pshufd(dest, dest, 0xfc);
         }
 
+        [Conditional("DEBUG")]
         private static void ValidateUnOp(Operand dest, Operand source)
         {
-#if DEBUG
             EnsureSameReg (dest, source);
             EnsureSameType(dest, source);
-#endif
         }
 
+        [Conditional("DEBUG")]
         private static void ValidateBinOp(Operand dest, Operand src1, Operand src2)
         {
-#if DEBUG
             EnsureSameReg (dest, src1);
             EnsureSameType(dest, src1, src2);
-#endif
         }
 
+        [Conditional("DEBUG")]
         private static void ValidateShift(Operand dest, Operand src1, Operand src2)
         {
-#if DEBUG
             EnsureSameReg (dest, src1);
             EnsureSameType(dest, src1);
 
             Debug.Assert(dest.Type.IsInteger() && src2.Type == OperandType.I32);
-#endif
         }
 
         private static void EnsureSameReg(Operand op1, Operand op2)
@@ -1595,7 +1592,7 @@ namespace ARMeilleure.CodeGen.X86
 
                 context.Assembler.Push(Register((X86Register)bit));
 
-                pushEntries.Add(new UnwindPushEntry(bit, RegisterType.Integer, context.StreamOffset));
+                pushEntries.Add(new UnwindPushEntry(UnwindPseudoOp.PushReg, context.StreamOffset, regIndex: bit));
 
                 mask &= ~(1 << bit);
             }
@@ -1612,6 +1609,8 @@ namespace ARMeilleure.CodeGen.X86
             if (reservedStackSize != 0)
             {
                 context.Assembler.Sub(rsp, Const(reservedStackSize), OperandType.I64);
+
+                pushEntries.Add(new UnwindPushEntry(UnwindPseudoOp.AllocStack, context.StreamOffset, stackOffsetOrAllocSize: reservedStackSize));
             }
 
             int offset = reservedStackSize;
@@ -1628,12 +1627,12 @@ namespace ARMeilleure.CodeGen.X86
 
                 context.Assembler.Movdqu(memOp, Xmm((X86Register)bit));
 
-                pushEntries.Add(new UnwindPushEntry(bit, RegisterType.Vector, context.StreamOffset));
+                pushEntries.Add(new UnwindPushEntry(UnwindPseudoOp.SaveXmm128, context.StreamOffset, bit, offset));
 
                 mask &= ~(1 << bit);
             }
 
-            return new UnwindInfo(pushEntries.ToArray(), context.StreamOffset, reservedStackSize);
+            return new UnwindInfo(pushEntries.ToArray(), context.StreamOffset);
         }
 
         private static void WriteEpilogue(CodeGenContext context)

ARMeilleure/Translation/JitUnwindWindows.cs

@@ -1,12 +1,15 @@
-using ARMeilleure.IntermediateRepresentation;
+// https://github.com/MicrosoftDocs/cpp-docs/blob/master/docs/build/exception-handling-x64.md
+
+using ARMeilleure.CodeGen.Unwinding;
 using System;
+using System.Diagnostics;
 using System.Runtime.InteropServices;
 
 namespace ARMeilleure.Translation
 {
     static class JitUnwindWindows
     {
-        private const int MaxUnwindCodesArraySize = 9 + 10 * 2 + 3;
+        private const int MaxUnwindCodesArraySize = 32; // Must be an even value.
 
         private struct RuntimeFunction
         {
@@ -41,12 +44,12 @@ namespace ARMeilleure.Translation
 
         [DllImport("kernel32.dll", CharSet = CharSet.Unicode)]
         private static unsafe extern bool RtlInstallFunctionTableCallback(
-            ulong                      tableIdentifier,
-            ulong                      baseAddress,
-            uint                       length,
+            ulong tableIdentifier,
+            ulong baseAddress,
+            uint length,
             GetRuntimeFunctionCallback callback,
-            IntPtr                     context,
-            string                     outOfProcessCallbackDll);
+            IntPtr context,
+            string outOfProcessCallbackDll);
 
         private static GetRuntimeFunctionCallback _getRuntimeFunctionCallback;
 
@@ -93,59 +96,80 @@ namespace ARMeilleure.Translation
 
             if (!JitCache.TryFind(offset, out JitCacheEntry funcEntry))
             {
-                // Not found.
-                return null;
+                return null; // Not found.
             }
 
             var unwindInfo = funcEntry.UnwindInfo;
 
             int codeIndex = 0;
 
-            int spOffset = unwindInfo.FixedAllocSize;
-
-            foreach (var entry in unwindInfo.PushEntries)
-            {
-                if (entry.Type == RegisterType.Vector)
-                {
-                    spOffset -= 16;
-                }
-            }
-
             for (int index = unwindInfo.PushEntries.Length - 1; index >= 0; index--)
             {
                 var entry = unwindInfo.PushEntries[index];
 
-                if (entry.Type == RegisterType.Vector)
+                switch (entry.PseudoOp)
                 {
-                    ushort uwop = PackUwop(UnwindOperation.SaveXmm128, entry.StreamEndOffset, entry.Index);
-
-                    _unwindInfo->UnwindCodes[codeIndex++] = uwop;
-                    _unwindInfo->UnwindCodes[codeIndex++] = (ushort)spOffset;
-
-                    spOffset += 16;
+                    case UnwindPseudoOp.SaveXmm128:
+                    {
+                        int stackOffset = entry.StackOffsetOrAllocSize;
+
+                        Debug.Assert(stackOffset % 16 == 0);
+
+                        if (stackOffset <= 0xFFFF0)
+                        {
+                            _unwindInfo->UnwindCodes[codeIndex++] = PackUnwindOp(UnwindOperation.SaveXmm128, entry.PrologOffset, entry.RegIndex);
+                            _unwindInfo->UnwindCodes[codeIndex++] = (ushort)(stackOffset / 16);
+                        }
+                        else
+                        {
+                            _unwindInfo->UnwindCodes[codeIndex++] = PackUnwindOp(UnwindOperation.SaveXmm128Far, entry.PrologOffset, entry.RegIndex);
+                            _unwindInfo->UnwindCodes[codeIndex++] = (ushort)(stackOffset >> 0);
+                            _unwindInfo->UnwindCodes[codeIndex++] = (ushort)(stackOffset >> 16);
+                        }
+
+                        break;
+                    }
+
+                    case UnwindPseudoOp.AllocStack:
+                    {
+                        int allocSize = entry.StackOffsetOrAllocSize;
+
+                        Debug.Assert(allocSize % 8 == 0);
+
+                        if (allocSize <= 128)
+                        {
+                            _unwindInfo->UnwindCodes[codeIndex++] = PackUnwindOp(UnwindOperation.AllocSmall, entry.PrologOffset, (allocSize / 8) - 1);
+                        }
+                        else if (allocSize <= 0x7FFF8)
+                        {
+                            _unwindInfo->UnwindCodes[codeIndex++] = PackUnwindOp(UnwindOperation.AllocLarge, entry.PrologOffset, 0);
+                            _unwindInfo->UnwindCodes[codeIndex++] = (ushort)(allocSize / 8);
+                        }
+                        else
+                        {
+                            _unwindInfo->UnwindCodes[codeIndex++] = PackUnwindOp(UnwindOperation.AllocLarge, entry.PrologOffset, 1);
+                            _unwindInfo->UnwindCodes[codeIndex++] = (ushort)(allocSize >> 0);
+                            _unwindInfo->UnwindCodes[codeIndex++] = (ushort)(allocSize >> 16);
+                        }
+
+                        break;
+                    }
+
+                    case UnwindPseudoOp.PushReg:
+                    {
+                        _unwindInfo->UnwindCodes[codeIndex++] = PackUnwindOp(UnwindOperation.PushNonvol, entry.PrologOffset, entry.RegIndex);
+
+                        break;
+                    }
+
+                    default: throw new NotImplementedException($"({nameof(entry.PseudoOp)} = {entry.PseudoOp})");
                 }
             }
 
-            _unwindInfo->UnwindCodes[0] = PackUwop(UnwindOperation.AllocLarge, unwindInfo.PrologueSize, 1);
-            _unwindInfo->UnwindCodes[1] = (ushort)(unwindInfo.FixedAllocSize >> 0);
-            _unwindInfo->UnwindCodes[2] = (ushort)(unwindInfo.FixedAllocSize >> 16);
-
-            codeIndex += 3;
-
-            for (int index = unwindInfo.PushEntries.Length - 1; index >= 0; index--)
-            {
-                var entry = unwindInfo.PushEntries[index];
-
-                if (entry.Type == RegisterType.Integer)
-                {
-                    ushort uwop = PackUwop(UnwindOperation.PushNonvol, entry.StreamEndOffset, entry.Index);
-
-                    _unwindInfo->UnwindCodes[codeIndex++] = uwop;
-                }
-            }
+            Debug.Assert(codeIndex <= MaxUnwindCodesArraySize);
 
-            _unwindInfo->VersionAndFlags    = 1;
-            _unwindInfo->SizeOfProlog       = (byte)unwindInfo.PrologueSize;
+            _unwindInfo->VersionAndFlags    = 1; // Flags: The function has no handler.
+            _unwindInfo->SizeOfProlog       = (byte)unwindInfo.PrologSize;
             _unwindInfo->CountOfUnwindCodes = (byte)codeIndex;
             _unwindInfo->FrameRegister      = 0;
 
@@ -156,9 +180,9 @@ namespace ARMeilleure.Translation
             return _runtimeFunction;
         }
 
-        private static ushort PackUwop(UnwindOperation uwop, int prologOffset, int opInfo)
+        private static ushort PackUnwindOp(UnwindOperation op, int prologOffset, int opInfo)
         {
-            return (ushort)(prologOffset | ((int)uwop << 8) | (opInfo << 12));
+            return (ushort)(prologOffset | ((int)op << 8) | (opInfo << 12));
         }
     }
 }