Propagate graph information in more directions

leanprover · Feb 25, 2025 · 0430515 · 0430515
1 parent 1a52635
commit 0430515
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/src/Lean/Elab/PreDefinition/Mutual.lean b/src/Lean/Elab/PreDefinition/Mutual.lean
@@ -91,12 +91,19 @@ partial def Info.setCallerParam (calleeIdx argIdx callerIdx paramIdx : Nat) (inf
       -- Set the new entry
       let info := info.modify calleeIdx (·.modify argIdx (·.map (·.set! callerIdx (some paramIdx))))
       Id.run do
-        -- Propagate
+        -- Propagate information about the caller
         let mut info : Info := info
         if let some callerParamInfo := info[callerIdx]![paramIdx]! then
           for h : otherFunIdx in [:callerParamInfo.size] do
             if let some otherParamIdx := callerParamInfo[otherFunIdx] then
               info := info.setCallerParam calleeIdx argIdx otherFunIdx otherParamIdx
+        -- Propagate information about the callee
+        for otherFunIdx in [:info.size] do
+          for otherArgIdx in [:info[otherFunIdx]!.size] do
+            if let some otherArgsInfo := info[otherFunIdx]![otherArgIdx]! then
+              if let some paramIdx' := otherArgsInfo[calleeIdx]! then
+                info := info.setCallerParam otherFunIdx otherArgIdx callerIdx paramIdx
+
         return info
     else
       -- Param not fixed, so argument isn't either