From 99dd60d8886155bb60768c8ef5a5a3225ad95ba4 Mon Sep 17 00:00:00 2001
From: Jakob Botsch Nielsen <Jakob.botsch.nielsen@gmail.com>
Date: Fri, 26 Apr 2024 10:00:38 +0200
Subject: [PATCH] JIT: Check more invariants on ABI info and fix some arm32
 bugs (#101372)

- Implicit byref parameters should be passed as a single pointer sized segment
- Segments should not point outside the local size
- Segments should not overlap each other
- Segments should have size > 0
- Segments should be ordered by offset
- Fix a bug in the arm32 classifier when structs are split
- Fix a bug in the arm32 classifier for odd-sized structs with 8 byte alignment.
  For example
```csharp
[StructLayout(LayoutKind.Sequential, Size = 12)]
struct S
{
    public double X;
    public float Y;
}
```
would be considered to take 4 slots before.
- Fix a bug in the Swift classifier that would cause the tail segments to have
  an out-of-bounds size
---
 src/coreclr/jit/abi.cpp           |  3 +++
 src/coreclr/jit/codegencommon.cpp |  4 ++++
 src/coreclr/jit/lclvars.cpp       | 36 +++++++++++++++++++++++++++++++
 src/coreclr/jit/targetarm.cpp     | 10 ++++-----
 4 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/src/coreclr/jit/abi.cpp b/src/coreclr/jit/abi.cpp
index fd899b899546b..c6dbd8aa9202f 100644
--- a/src/coreclr/jit/abi.cpp
+++ b/src/coreclr/jit/abi.cpp
@@ -415,6 +415,9 @@ ABIPassingInformation SwiftABIClassifier::Classify(Compiler*    comp,
             {
                 ABIPassingSegment newSegment = elemInfo.Segments[j];
                 newSegment.Offset += lowering->offsets[i];
+                // Adjust the tail size if necessary; the lowered sequence can
+                // pass the tail as a larger type than the tail size.
+                newSegment.Size = min(newSegment.Size, structLayout->GetSize() - newSegment.Offset);
                 segments.Push(newSegment);
             }
         }
diff --git a/src/coreclr/jit/codegencommon.cpp b/src/coreclr/jit/codegencommon.cpp
index c8211a7e73144..9f3120d274d57 100644
--- a/src/coreclr/jit/codegencommon.cpp
+++ b/src/coreclr/jit/codegencommon.cpp
@@ -4185,9 +4185,13 @@ void CodeGen::genHomeSwiftStructParameters(bool handleStack)
                     case 2:
                         loadType = TYP_USHORT;
                         break;
+                    case 3:
                     case 4:
                         loadType = TYP_INT;
                         break;
+                    case 5:
+                    case 6:
+                    case 7:
                     case 8:
                         loadType = TYP_LONG;
                         break;
diff --git a/src/coreclr/jit/lclvars.cpp b/src/coreclr/jit/lclvars.cpp
index e6645e1f03d8f..97986eefee4b2 100644
--- a/src/coreclr/jit/lclvars.cpp
+++ b/src/coreclr/jit/lclvars.cpp
@@ -1885,6 +1885,42 @@ void Compiler::lvaClassifyParameterABI()
             }
         }
     }
+
+    for (unsigned lclNum = 0; lclNum < info.compArgsCount; lclNum++)
+    {
+        const ABIPassingInformation& abiInfo = lvaGetParameterABIInfo(lclNum);
+
+        if (lvaIsImplicitByRefLocal(lclNum))
+        {
+            assert((abiInfo.NumSegments == 1) && (abiInfo.Segments[0].Size == TARGET_POINTER_SIZE));
+        }
+        else
+        {
+            for (unsigned i = 0; i < abiInfo.NumSegments; i++)
+            {
+                const ABIPassingSegment& segment = abiInfo.Segments[i];
+                assert(segment.Size > 0);
+                assert(segment.Offset + segment.Size <= lvaLclExactSize(lclNum));
+
+                if (i > 0)
+                {
+                    assert(segment.Offset > abiInfo.Segments[i - 1].Offset);
+                }
+
+                for (unsigned j = 0; j < abiInfo.NumSegments; j++)
+                {
+                    if (i == j)
+                    {
+                        continue;
+                    }
+
+                    const ABIPassingSegment& otherSegment = abiInfo.Segments[j];
+                    assert((segment.Offset + segment.Size <= otherSegment.Offset) ||
+                           (segment.Offset >= otherSegment.Offset + otherSegment.Size));
+                }
+            }
+        }
+    }
 #endif // DEBUG
 }
 
diff --git a/src/coreclr/jit/targetarm.cpp b/src/coreclr/jit/targetarm.cpp
index 037578fa67b85..675fd04230d53 100644
--- a/src/coreclr/jit/targetarm.cpp
+++ b/src/coreclr/jit/targetarm.cpp
@@ -85,11 +85,11 @@ ABIPassingInformation Arm32Classifier::Classify(Compiler*    comp,
         m_nextIntReg = roundUp(m_nextIntReg, 2);
     }
 
-    unsigned size        = type == TYP_STRUCT ? structLayout->GetSize() : genTypeSize(type);
-    unsigned alignedSize = roundUp(size, alignment);
+    unsigned size     = type == TYP_STRUCT ? structLayout->GetSize() : genTypeSize(type);
+    unsigned numSlots = (size + 3) / 4;
 
-    unsigned numInRegs  = min(alignedSize / 4, 4 - m_nextIntReg);
-    bool     anyOnStack = numInRegs < (alignedSize / 4);
+    unsigned numInRegs  = min(numSlots, 4 - m_nextIntReg);
+    bool     anyOnStack = numInRegs < numSlots;
 
     // If we already passed anything on stack (due to float args) then we
     // cannot split an arg.
@@ -116,7 +116,7 @@ ABIPassingInformation Arm32Classifier::Classify(Compiler*    comp,
     {
         m_stackArgSize           = roundUp(m_stackArgSize, alignment);
         unsigned stackSize       = size - (numInRegs * 4);
-        info.Segments[numInRegs] = ABIPassingSegment::OnStack(m_stackArgSize, 0, stackSize);
+        info.Segments[numInRegs] = ABIPassingSegment::OnStack(m_stackArgSize, numInRegs * 4, stackSize);
         m_stackArgSize += roundUp(stackSize, 4);
 
         // As soon as any int arg goes on stack we cannot put anything else in