1 files changed, 22 insertions, 4 deletions
diff --git a/ChocolArm64/Translation/ILOpCodeStoreState.cs b/ChocolArm64/Translation/ILOpCodeStoreState.cs
index 458e9eda..a587dbfe 100644
--- a/ChocolArm64/Translation/ILOpCodeStoreState.cs
+++ b/ChocolArm64/Translation/ILOpCodeStoreState.cs
@@ -7,15 +7,33 @@ namespace ChocolArm64.Translation
     {
         private ILBlock _block;
 
-        public ILOpCodeStoreState(ILBlock block)
+        private TranslatedSub _callSub;
+
+        public ILOpCodeStoreState(ILBlock block, TranslatedSub callSub = null)
         {
-            _block = block;
+            _block   = block;
+            _callSub = callSub;
         }
 
         public void Emit(ILMethodBuilder context)
         {
-            long intOutputs = context.LocalAlloc.GetIntOutputs(_block);
-            long vecOutputs = context.LocalAlloc.GetVecOutputs(_block);
+            long intOutputs = context.RegUsage.GetIntOutputs(_block);
+            long vecOutputs = context.RegUsage.GetVecOutputs(_block);
+
+            if (Optimizations.AssumeStrictAbiCompliance && context.IsSubComplete)
+            {
+                intOutputs = RegisterUsage.ClearCallerSavedIntRegs(intOutputs, context.IsAarch64);
+                vecOutputs = RegisterUsage.ClearCallerSavedVecRegs(vecOutputs, context.IsAarch64);
+            }
+
+            if (_callSub != null)
+            {
+                //Those register are assigned on the callee function, without
+                //reading it's value first. We don't need to write them because
+                //they are not going to be read on the callee.
+                intOutputs &= ~_callSub.IntNiRegsMask;
+                vecOutputs &= ~_callSub.VecNiRegsMask;
+            }
 
             StoreLocals(context, intOutputs, RegisterType.Int);
             StoreLocals(context, vecOutputs, RegisterType.Vector);