refactor: Hyps optimisations (experiment)

Iris/Iris/ProofMode/Display.lean

@@ -61,7 +61,7 @@ where
       else
         `(irisHyp| ∗$(mkIdent name') : $(← unpackIprop (← delab ty)))
       pure (map.insert name idx, acc.push stx)
-    | .sep _ _ _ _ lhs rhs => delabHypotheses lhs (← delabHypotheses rhs acc)
+    | .sep _ _ _ _ _ lhs rhs => delabHypotheses lhs (← delabHypotheses rhs acc)
 
 @[delab app.Iris.ProofMode.HypMarker]
 def delabHypMarker : Delab := do unpackIprop (← withAppArg delab)

Iris/Iris/ProofMode/Expr.lean

@@ -128,14 +128,20 @@ This means that the ivar correctly caches whether it refers to a persistent hypo
 -/
 inductive Hyps {prop : Q(Type u)} (bi : Q(BI $prop)) : (e : Q($prop)) → Type where
   | emp (_ : $e =Q emp) : Hyps bi e
-  | sep (tm elhs erhs : Q($prop)) (_ : $e =Q iprop($elhs ∗ $erhs))
+  | sep (tm elhs erhs : Q($prop)) (size : Nat) (_ : $e =Q iprop($elhs ∗ $erhs))
         (lhs : Hyps bi elhs) (rhs : Hyps bi erhs) : Hyps bi e
   | hyp (tm : Q($prop)) (name : Name) (ivar : IVarId) (p : Q(Bool)) (ty : Q($prop))
         (_ : $e =Q iprop(□?$p $ty)) : Hyps bi e
 deriving Repr
 
 instance : Inhabited (Hyps bi s) := ⟨.emp ⟨⟩⟩
 
+@[inline]
+def Hyps.size : @Hyps u prop bi s → Nat
+  | .emp _ => 0
+  | .hyp .. => 1
+  | .sep _ _ _ n .. => n
+
 def Hyps.tm : @Hyps _ prop bi s → Q($prop)
   | .emp _ => s
   | .sep tm .. | .hyp tm .. => tm
@@ -144,7 +150,7 @@ def Hyps.mkEmp {prop : Q(Type u)} (bi : Q(BI $prop)) (e := q(BI.emp : $prop)) :
 
 def Hyps.mkSep {prop : Q(Type u)} {bi : Q(BI $prop)} {elhs erhs}
     (lhs : Hyps bi elhs) (rhs : Hyps bi erhs) (e := q(BI.sep $elhs $erhs)) : Hyps bi e :=
-  .sep q(BI.sep $(lhs.tm) $(rhs.tm) : $prop) elhs erhs ⟨⟩ lhs rhs
+  .sep q(BI.sep $(lhs.tm) $(rhs.tm) : $prop) elhs erhs (lhs.size + rhs.size) ⟨⟩ lhs rhs
 
 def mkIntuitionisticIf {prop : Q(Type u)} (_bi : Q(BI $prop))
     (p : Q(Bool)) (e : Q($prop)) : {A : Q($prop) // $A =Q iprop(□?$p $e)} :=
@@ -156,20 +162,78 @@ def Hyps.mkHyp {prop : Q(Type u)} (bi : Q(BI $prop))
     (name : Name) (ivar : IVarId) (p : Q(Bool)) (ty : Q($prop)) (e := q(iprop(□?$p $ty))) : Hyps bi e :=
   .hyp (mkIntuitionisticIf bi p (mkNameAnnotation name ivar ty)) name ivar p ty ⟨⟩
 
+/-- A `Hyps` tree together with a proof that it represents `e`. -/
+structure HypsEq {prop : Q(Type u)} (bi : Q(BI $prop)) (e : Q($prop)) where
+  (e' : Q($prop)) (hyps : Hyps bi e') (pf : Q($e ⊣⊢ $e'))
+
+def wbDelta : Nat := 3
+def wbGamma : Nat := 2
+
+/-- Rebuild `lhs ∗ rhs` when `rhs` may be one element too heavy. -/
+private def Hyps.balanceR {u} {prop : Q(Type u)} {bi : Q(BI $prop)}
+    {elhs erhs : Q($prop)} (lhs : Hyps bi elhs) (rhs : Hyps bi erhs) :
+    HypsEq bi q(iprop($elhs ∗ $erhs)) :=
+  let ln := lhs.size; let rn := rhs.size
+  if ln + rn ≤ 1 || rn ≤ wbDelta * ln then ⟨_, .mkSep lhs rhs, q(.rfl)⟩ else
+  match rhs with
+  | .sep _ erl err _ _ rl rr =>
+    have : $erhs =Q iprop($erl ∗ $err) := ⟨⟩
+    if rl.size < wbGamma * rr.size then
+      -- single:  A ∗ (B ∗ C)  ⇝  (A ∗ B) ∗ C
+      ⟨_, .mkSep (.mkSep lhs rl) rr, q(sep_assoc.symm)⟩
+    else match rl with
+      | .sep _ erll erlr _ _ rll rlr =>
+        have : $erl =Q iprop($erll ∗ $erlr) := ⟨⟩
+        -- double:  A ∗ ((B ∗ C) ∗ D)  ⇝  (A ∗ B) ∗ (C ∗ D)
+        ⟨_, .mkSep (.mkSep lhs rll) (.mkSep rlr rr), q((sep_congr_right sep_assoc).trans sep_assoc.symm)⟩
+      | _ => ⟨_, .mkSep lhs rhs, q(.rfl)⟩
+  | _ => ⟨_, .mkSep lhs rhs, q(.rfl)⟩
+
+/-- Rebuild `lhs ∗ rhs` when `lhs` may be one element too heavy. -/
+private def Hyps.balanceL {u} {prop : Q(Type u)} {bi : Q(BI $prop)}
+    {elhs erhs : Q($prop)} (lhs : Hyps bi elhs) (rhs : Hyps bi erhs) :
+    HypsEq bi q(iprop($elhs ∗ $erhs)) :=
+  let ln := lhs.size; let rn := rhs.size
+  if ln + rn ≤ 1 || ln ≤ wbDelta * rn then ⟨_, .mkSep lhs rhs, q(.rfl)⟩ else
+  match lhs with
+  | .sep _ ell elr _ _ ll lr =>
+    have : $elhs =Q iprop($ell ∗ $elr) := ⟨⟩
+    if lr.size < wbGamma * ll.size then
+      -- single:  (A ∗ B) ∗ C  ⇝  A ∗ (B ∗ C)
+      ⟨_, .mkSep ll (.mkSep lr rhs), q(sep_assoc)⟩
+    else match lr with
+      | .sep _ elrl elrr _ _ lrl lrr =>
+        have : $elr =Q iprop($elrl ∗ $elrr) := ⟨⟩
+        -- double:  (A ∗ (B ∗ C)) ∗ D  ⇝  (A ∗ B) ∗ (C ∗ D)
+        ⟨_, .mkSep (.mkSep ll lrl) (.mkSep lrr rhs), q((sep_congr_left sep_assoc.symm).trans sep_assoc)⟩
+      | _ => ⟨_, .mkSep lhs rhs, q(.rfl)⟩
+  | _ => ⟨_, .mkSep lhs rhs, q(.rfl)⟩
+
+/-- Append `leaf` at the right end of `hyps`, rebalancing on the way back up. -/
+private def Hyps.snoc {u} {prop : Q(Type u)} {bi : Q(BI $prop)}
+    {e ehyp : Q($prop)} (hyps : Hyps bi e) (leaf : Hyps bi ehyp) :
+    HypsEq bi q(iprop($e ∗ $ehyp)) :=
+  match hyps with
+  | .emp _  => ⟨_, leaf, q(emp_sep)⟩
+  | .hyp .. => ⟨_, .mkSep hyps leaf, q(.rfl)⟩
+  | .sep _ elhs erhs _ _ lhs rhs =>
+    have : $e =Q iprop($elhs ∗ $erhs) := ⟨⟩
+    let ⟨_, rhs', pfR⟩ := rhs.snoc leaf              -- erhs ∗ ehyp ⊣⊢ erhs'
+    let ⟨_, hyps', pfB⟩ := Hyps.balanceR lhs rhs'    -- elhs ∗ erhs' ⊣⊢ e''
+    ⟨_, hyps', q(sep_assoc.trans ((sep_congr_right $pfR).trans $pfB))⟩
+
 def Hyps.add {prop : Q(Type u)} (bi : Q(BI $prop))
-    (name : Name) (ivar : IVarId) (p : Q(Bool)) (ty : Q($prop)) {e} (h : Hyps bi e)
-    : (e' : Q($prop)) × Hyps bi e' × Q(iprop($e ∗ □?$p $ty ⊣⊢ $e')) :=
-  match h with
-  -- Adding a hypothesis to `emp` creates a `.hyp` node instead of a `.sep` node
-  | .emp _ => ⟨_, .mkHyp bi name ivar p ty, q(emp_sep)⟩
-  | _ => ⟨_, .mkSep h (.mkHyp bi name ivar p ty), q(.rfl)⟩
+    (name : Name) (ivar : IVarId) (p : Q(Bool)) (ty : Q($prop)) {e} (h : Hyps bi e) :
+    (e' : Q($prop)) × Hyps bi e' × Q(iprop($e ∗ □?$p $ty ⊣⊢ $e')) :=
+  let ⟨e', hyps', pf⟩ := h.snoc <| Hyps.mkHyp bi name ivar p ty
+  ⟨e', hyps', pf⟩
 
 partial def parseHyps? {prop : Q(Type u)} (bi : Q(BI $prop)) (expr : Expr) :
     Option ((s : Q($prop)) × Hyps bi s) := do
   if let some #[_, _, P, Q] := appM? expr ``sep then
     let ⟨elhs, lhs⟩ ← parseHyps? bi P
     let ⟨erhs, rhs⟩ ← parseHyps? bi Q
-    some ⟨q(BI.sep $elhs $erhs), .sep expr elhs erhs ⟨⟩ lhs rhs⟩
+    some ⟨q(BI.sep $elhs $erhs), .sep expr elhs erhs (lhs.size + rhs.size) ⟨⟩ lhs rhs⟩
   else if expr.isAppOfArity ``emp 2 then
     some ⟨expr, .emp ⟨⟩⟩
   else if let some #[_, _, P] := appM? expr ``intuitionistically then
@@ -183,7 +247,7 @@ partial def Hyps.find? {u prop bi} (name : Name) :
     ∀ {s}, @Hyps u prop bi s → Option (IVarId × Q($prop))
   | _, .emp _ => none
   | _, .hyp _ name' ivar _ ty _ => if name == name' then (ivar, ty) else none
-  | _, .sep _ _ _ _ lhs rhs => rhs.find? name <|> lhs.find? name
+  | _, .sep _ _ _ _ _ lhs rhs => rhs.find? name <|> lhs.find? name
 
 partial def Hyps.findM? [Monad m] {prop : Q(Type u)} {bi : Q(BI $prop)}
     (p : Name → IVarId → Q(Bool) → Q($prop) → m Bool) :
@@ -194,7 +258,7 @@ partial def Hyps.findM? [Monad m] {prop : Q(Type u)} {bi : Q(BI $prop)}
       return some (name, ivar, bp, ty)
     else
       return none
-  | _, .sep _ _ _ _ lhs rhs => do
+  | _, .sep _ _ _ _ _ lhs rhs => do
     match ← rhs.findM? p with
     | some res => return some res
     | none => lhs.findM? p
@@ -203,22 +267,24 @@ partial def Hyps.getDecl? {u prop bi} (ivar : IVarId) {s}:
     @Hyps u prop bi s → Option (Name × IVarId × Q(Bool) × Q($prop))
   | .emp _ => none
   | .hyp _ name ivar' p ty _ => if ivar == ivar' then (name, ivar, p, ty) else none
-  | .sep _ _ _ _ lhs rhs => rhs.getDecl? ivar <|> lhs.getDecl? ivar
+  | .sep _ _ _ _ _ lhs rhs => rhs.getDecl? ivar <|> lhs.getDecl? ivar
 
 def Hyps.getUserName? {u prop bi} (ivar : IVarId) (h : @Hyps u prop bi s) : Option Name :=
   h.getDecl? ivar |>.map (·.1)
 
-partial def Hyps.spatialIVarIds {u prop bi} :
-    ∀ {s}, @Hyps u prop bi s → List IVarId
-  | _, .emp _ => []
-  | _, .hyp _ _ ivar p _ _ => if isTrue p then [] else [ivar]
-  | _, .sep _ _ _ _ lhs rhs => lhs.spatialIVarIds ++ rhs.spatialIVarIds
+@[specialize]
+def Hyps.foldrLeaves {u} {prop : Q(Type u)} {bi : Q(BI $prop)} {α}
+    (f : Name → IVarId → Q(Bool) → Q($prop) → α → α) :
+    ∀ {s}, @Hyps u prop bi s → α → α
+  | _, .emp _,               acc => acc
+  | _, .hyp _ n ivar p ty _, acc => f n ivar p ty acc
+  | _, .sep _ _ _ _ _ lhs rhs, acc => foldrLeaves f lhs (foldrLeaves f rhs acc)
+
+def Hyps.spatialIVarIds {u prop bi} {s} (hyps : @Hyps u prop bi s) : List IVarId :=
+  hyps.foldrLeaves (fun _ ivar p _ acc => if isTrue p then acc else ivar :: acc) []
 
-partial def Hyps.intuitionisticIVarIds {u prop bi} :
-    ∀ {s}, @Hyps u prop bi s → List IVarId
-  | _, .emp _ => []
-  | _, .hyp _ _ ivar p _ _ => if isTrue p then [ivar] else []
-  | _, .sep _ _ _ _ lhs rhs => lhs.intuitionisticIVarIds ++ rhs.intuitionisticIVarIds
+def Hyps.intuitionisticIVarIds {u prop bi} {s} (hyps : @Hyps u prop bi s) : List IVarId :=
+  hyps.foldrLeaves (fun _ ivar p _ acc => if isTrue p then ivar :: acc else acc) []
 
 /--
   Given any hypotheses `hyps` representing `e`, filter in all spatial hypotheses
@@ -244,15 +310,15 @@ def Hyps.buildAccuProof {prop : Q(Type u)} {bi : Q(BI $prop)} {e}
           let pf : Q($e' ⊢ iprop(emp)) := pf
           ⟨ty, q((sep_mono_right $pf).trans sep_emp.mp)⟩
         else ⟨q(iprop($ty ∗ $spatialProps)), q(sep_mono_right $pf)⟩
-    | .sep _ _ _ _ lhs rhs =>
+    | .sep _ _ _ _ _ lhs rhs =>
       let ⟨spatialPropsR, pfR⟩ := buildAccuProofAux rhs spatialProps pf
       let ⟨spatialPropsLR, pfLR⟩ := buildAccuProofAux lhs spatialPropsR pfR
       ⟨q($spatialPropsLR), q(sep_assoc.mp.trans $pfLR)⟩
 
 variable (oldIVar : IVarId) (new : Name) {prop : Q(Type u)} {bi : Q(BI $prop)} in
 def Hyps.rename : ∀ {e}, Hyps bi e → Option (Hyps bi e)
   | _, .emp _ => none
-  | _, .sep _ _ _ _ lhs rhs =>
+  | _, .sep _ _ _ _ _ lhs rhs =>
     match rhs.rename with
     | some rhs' => some (.mkSep lhs rhs' _)
     | none => match lhs.rename with
@@ -266,7 +332,7 @@ def Hyps.select (ty : Expr) : ∀ {s}, @Hyps u prop bi s → MetaM (IVarId × Q(
   | _, .hyp _ _ ivar p ty' _ => do
     let .true ← isDefEq ty ty' | failure
     pure (ivar, p, ty')
-  | _, .sep _ _ _ _ lhs rhs => try Hyps.select ty rhs catch _ => Hyps.select ty lhs
+  | _, .sep _ _ _ _ _ lhs rhs => try Hyps.select ty rhs catch _ => Hyps.select ty lhs
 
 
 theorem intuitionistically_sep_dup [BI PROP] {P : PROP} : □ P ⊣⊢ □ P ∗ □ P :=
@@ -310,7 +376,7 @@ def Hyps.splitCore : ∀ {e}, Hyps bi e → SplitResult bi e
       have : $ehyp =Q iprop(□ $ty) := ⟨⟩
       .split h h q(intuitionistically_sep_dup)
     | .inr _ => if toRight name ivar then .right else .left
-  | _, .sep _ _ _ _ lhs rhs =>
+  | _, .sep _ _ _ _ _ lhs rhs =>
     let resl := lhs.splitCore
     let resr := rhs.splitCore
     match resl, resr with
@@ -372,19 +438,21 @@ def Hyps.removeCore : ∀ {e}, Hyps bi e → m (RemoveHypCore bi e α)
       | _, _ => return .one a ty p ⟨⟩
     else
       return .none
-  | _, .sep _ elhs erhs _ lhs rhs => do
+  | _, .sep _ elhs erhs _ _ lhs rhs => do
     match ← rhs.removeCore with
     | .one a out' p h =>
+      -- the whole right child was a leaf and is gone; nothing to rebalance here
       return .main a ⟨elhs, lhs, erhs, out', p, h, q(.rfl)⟩
     | .main a ⟨_, rhs', out, out', p, h, pf⟩ =>
-      let hyps' := .mkSep lhs rhs'
-      return .main a ⟨_, hyps', out, out', p, h, q(remove_r $pf)⟩
+      -- the right child shrank ⇒ the left may now be too heavy
+      let ⟨_, hyps', bal⟩ := Hyps.balanceL lhs rhs'
+      return .main a ⟨_, hyps', out, out', p, h, q((remove_r $pf).trans (sep_congr_left $bal))⟩
     | .none => match ← lhs.removeCore with
       | .one a out' p h =>
         return .main a ⟨erhs, rhs, elhs, out', p, h, q(sep_comm)⟩
       | .main a ⟨_, lhs', out, out', p, h, pf⟩ =>
-        let hyps' := .mkSep lhs' rhs
-        return .main a ⟨_, hyps', out, out', p, h, q(remove_l $pf)⟩
+        let ⟨_, hyps', bal⟩ := Hyps.balanceR lhs' rhs
+        return .main a ⟨_, hyps', out, out', p, h, q((remove_l $pf).trans (sep_congr_left $bal))⟩
       | .none => pure .none
 
 def Hyps.removeG [Monad m] {prop : Q(Type u)} {bi : Q(BI $prop)} {e : Q(Prop)}
@@ -481,7 +549,7 @@ def Hyps.replaceCore : ∀ {e}, Hyps bi e → m (Option ((e' : Q($prop)) × Hyps
       let ⟨ty', pf⟩ ← repl name p ty
       return some ⟨_, .mkHyp bi name ivar p ty', q(replace_hyp $pf)⟩
     return none
-  | _, .sep _ _ _ _ lhs rhs => do
+  | _, .sep _ _ _ _ _ lhs rhs => do
     if let some ⟨_, lhs', pf⟩ ← lhs.replaceCore then
       return some ⟨_, .mkSep lhs' rhs, q(replace_hyp_sep_l $pf)⟩
     if let some ⟨_, rhs', pf⟩ ← rhs.replaceCore then
@@ -501,7 +569,7 @@ section dependency
 partial def Hyps.findDependencyOnFVar {prop : Q(Type u)} {bi : Q(BI $prop)}
     (fvarId : FVarId) : ∀ {e}, Hyps bi e → Option (Name × IVarId × Q(Bool) × Q($prop))
   | _, .emp _ => none
-  | _, .sep _ _ _ _ lhs rhs =>
+  | _, .sep _ _ _ _ _ lhs rhs =>
       lhs.findDependencyOnFVar fvarId <|> rhs.findDependencyOnFVar fvarId
   | _, .hyp _ name ivar p ty _ =>
       if (ty : Expr).containsFVar fvarId then some (name, ivar, p, ty)
@@ -621,7 +689,7 @@ def Hyps.buildIntuitionisticProof {u} {prop : Q(Type u)} {bi : Q(BI $prop)} {e}
     match matchBool p with
     | .inl _ => some q(intuitionistically_idem.mpr)
     | .inr _ => none
-  | .sep _ _ _ _ lhs rhs => do
+  | .sep _ _ _ _ _ lhs rhs => do
     let pfL ← buildIntuitionisticProof lhs
     let pfR ← buildIntuitionisticProof rhs
     some q((sep_mono $pfL $pfR).trans intuitionistically_sep_mpr)

Iris/Iris/ProofMode/Tactics/ModIntro.lean

@@ -137,7 +137,7 @@ where go {e}
       have heq : Q(@ModalityAction.id $prop1 = .id) := q(Eq.refl (ModalityAction.id))
       have heq : Q($(M).action $p = .id) := heq
       return ⟨_, .mkHyp bi1 name ivar p ty, q(modaction_id $M $heq)⟩
-  | .sep _ _ _ _ lhs rhs => do
+  | .sep _ _ _ _ _ lhs rhs => do
     let ⟨_, lhs', pflhs⟩ ← go iact sact lhs
     let ⟨_, rhs', pfrhs⟩ ← go iact sact rhs
     -- TODO: make pruning emp part of mkSep?

Iris/Iris/Tests/Tactics.lean

@@ -1074,7 +1074,7 @@ example [BI PROP] (P Q : PROP) : (□P ∗ Q) -∗ Q := by
 
 /-- Tests `ihave` not removing a destructed hyp -/
 example [BI PROP] [BIAffine PROP] (Q : PROP) :
-  □ (Q ∗ Q) ⊢ (□ (Q ∗ Q) ∗ □ Q) ∗ □ Q := by
+  □ (Q ∗ Q) ⊢ □ (Q ∗ Q) ∗ □ Q ∗ □ Q := by
   iintro #HQ
   ihave ⟨HQ, HQ2⟩ := HQ
   istop