perf: use RArray in simp_arith meta code (#6068 part 2)

This PR makes `simp_arith` use `RArray` for the context of the
reflection proofs, which scales better when there are many variables.

On our synthetic benchmark:
```
simp_arith1               instructions    -25.1% (-4892.6 σ)
```

No effect on mathlib, though, guess it’s not used much on large goals there:
http://speed.lean-fro.org/mathlib4/compare/873b982b-2038-462a-9b68-0c0fc457f90d/to/56e66691-2f1f-4947-a922-37b80680315d

src/Init/Data/Nat/Linear.lean

@@ -6,6 +6,7 @@ Authors: Leonardo de Moura
 prelude
 import Init.ByCases
 import Init.Data.Prod
+import Init.Data.RArray
 
 namespace Nat.Linear
 
@@ -15,7 +16,7 @@ namespace Nat.Linear
 
 abbrev Var := Nat
 
-abbrev Context := List Nat
+abbrev Context := Lean.RArray Nat
 
 /--
   When encoding polynomials. We use `fixedVar` for encoding numerals.
@@ -23,12 +24,7 @@ abbrev Context := List Nat
 def fixedVar := 100000000 -- Any big number should work here
 
 def Var.denote (ctx : Context) (v : Var) : Nat :=
-  bif v == fixedVar then 1 else go ctx v
-where
-  go : List Nat → Nat → Nat
-   | [],    _   => 0
-   | a::_,  0   => a
-   | _::as, i+1 => go as i
+  bif v == fixedVar then 1 else ctx.get v
 
 inductive Expr where
   | num  (v : Nat)

src/Lean/Meta/Tactic/LinearArith/Nat/Basic.lean

@@ -8,6 +8,7 @@ import Lean.Meta.Check
 import Lean.Meta.Offset
 import Lean.Meta.AppBuilder
 import Lean.Meta.KExprMap
+import Lean.Data.RArray
 
 namespace Lean.Meta.Linear.Nat
 
@@ -141,8 +142,11 @@ end ToLinear
 
 export ToLinear (toLinearCnstr? toLinearExpr)
 
-def toContextExpr (ctx : Array Expr) : MetaM Expr := do
-  mkListLit (mkConst ``Nat) ctx.toList
+def toContextExpr (ctx : Array Expr) : Expr :=
+  if h : 0 < ctx.size then
+    RArray.toExpr (mkConst ``Nat) id (RArray.ofArray ctx h)
+  else
+    RArray.toExpr (mkConst ``Nat) id (RArray.leaf (mkNatLit 0))
 
 def reflTrue : Expr :=
   mkApp2 (mkConst ``Eq.refl [levelOne]) (mkConst ``Bool) (mkConst ``Bool.true)

src/Lean/Meta/Tactic/LinearArith/Nat/Simp.lean

@@ -31,17 +31,17 @@ def simpCnstrPos? (e : Expr) : MetaM (Option (Expr × Expr)) := do
     let c₂ := c₁.norm
     if c₂.isUnsat then
       let r := mkConst ``False
-      let p := mkApp3 (mkConst ``Nat.Linear.ExprCnstr.eq_false_of_isUnsat) (← toContextExpr ctx) (toExpr c) reflTrue
+      let p := mkApp3 (mkConst ``Nat.Linear.ExprCnstr.eq_false_of_isUnsat) (toContextExpr ctx) (toExpr c) reflTrue
       return some (r, ← mkExpectedTypeHint p (← mkEq lhs r))
     else if c₂.isValid then
       let r := mkConst ``True
-      let p := mkApp3 (mkConst ``Nat.Linear.ExprCnstr.eq_true_of_isValid) (← toContextExpr ctx) (toExpr c) reflTrue
+      let p := mkApp3 (mkConst ``Nat.Linear.ExprCnstr.eq_true_of_isValid) (toContextExpr ctx) (toExpr c) reflTrue
       return some (r, ← mkExpectedTypeHint p (← mkEq lhs r))
     else
       let c₂ : LinearCnstr := c₂.toExpr
       let r ← c₂.toArith ctx
       if r != lhs then
-        let p := mkApp4 (mkConst ``Nat.Linear.ExprCnstr.eq_of_toNormPoly_eq) (← toContextExpr ctx) (toExpr c) (toExpr c₂) reflTrue
+        let p := mkApp4 (mkConst ``Nat.Linear.ExprCnstr.eq_of_toNormPoly_eq) (toContextExpr ctx) (toExpr c) (toExpr c₂) reflTrue
         return some (r, ← mkExpectedTypeHint p (← mkEq lhs r))
       else
         return none
@@ -81,7 +81,7 @@ def simpExpr? (e : Expr) : MetaM (Option (Expr × Expr)) := do
   if p'.length < p.length then
     -- We only return some if monomials were fused
     let e' : LinearExpr := p'.toExpr
-    let p := mkApp4 (mkConst ``Nat.Linear.Expr.eq_of_toNormPoly_eq) (← toContextExpr ctx) (toExpr e) (toExpr e') reflTrue
+    let p := mkApp4 (mkConst ``Nat.Linear.Expr.eq_of_toNormPoly_eq) (toContextExpr ctx) (toExpr e) (toExpr e') reflTrue
     let r ← e'.toArith ctx
     return some (r, p)
   else

tests/lean/run/linearByRefl.lean

@@ -1,13 +1,25 @@
+import Lean
+
 open Nat.Linear
 
+-- Convenient RArray literals
+elab tk:"#R[" ts:term,* "]" : term => do
+  let ts : Array Lean.Syntax := ts
+  let es ← ts.mapM fun stx => Lean.Elab.Term.elabTerm stx none
+  if h : 0 < es.size then
+    return (Lean.RArray.toExpr (← Lean.Meta.inferType es[0]!) id (Lean.RArray.ofArray es h))
+  else
+    throwErrorAt tk "RArray cannot be empty"
+
+
 example (x₁ x₂ x₃ : Nat) : (x₁ + x₂) + (x₂ + x₃) = x₃ + 2*x₂ + x₁ :=
-  Expr.eq_of_toNormPoly [x₁, x₂, x₃]
+  Expr.eq_of_toNormPoly #R[x₁, x₂, x₃]
    (Expr.add (Expr.add (Expr.var 0) (Expr.var 1)) (Expr.add (Expr.var 1) (Expr.var 2)))
    (Expr.add (Expr.add (Expr.var 2) (Expr.mulL 2 (Expr.var 1))) (Expr.var 0))
    rfl
 
 example (x₁ x₂ x₃ : Nat) : ((x₁ + x₂) + (x₂ + x₃) = x₃ + x₂) = (x₁ + x₂ = 0) :=
-  Expr.of_cancel_eq [x₁, x₂, x₃]
+  Expr.of_cancel_eq #R[x₁, x₂, x₃]
     (Expr.add (Expr.add (Expr.var 0) (Expr.var 1)) (Expr.add (Expr.var 1) (Expr.var 2)))
     (Expr.add (Expr.var 2) (Expr.var 1))
     (Expr.add (Expr.var 0) (Expr.var 1))
@@ -15,7 +27,7 @@ example (x₁ x₂ x₃ : Nat) : ((x₁ + x₂) + (x₂ + x₃) = x₃ + x₂) =
     rfl
 
 example (x₁ x₂ x₃ : Nat) : ((x₁ + x₂) + (x₂ + x₃) ≤ x₃ + x₂) = (x₁ + x₂ ≤ 0) :=
-  Expr.of_cancel_le [x₁, x₂, x₃]
+  Expr.of_cancel_le #R[x₁, x₂, x₃]
     (Expr.add (Expr.add (Expr.var 0) (Expr.var 1)) (Expr.add (Expr.var 1) (Expr.var 2)))
     (Expr.add (Expr.var 2) (Expr.var 1))
     (Expr.add (Expr.var 0) (Expr.var 1))
@@ -23,7 +35,7 @@ example (x₁ x₂ x₃ : Nat) : ((x₁ + x₂) + (x₂ + x₃) ≤ x₃ + x₂)
     rfl
 
 example (x₁ x₂ x₃ : Nat) : ((x₁ + x₂) + (x₂ + x₃) < x₃ + x₂) = (x₁ + x₂ < 0) :=
-  Expr.of_cancel_lt [x₁, x₂, x₃]
+  Expr.of_cancel_lt #R[x₁, x₂, x₃]
     (Expr.add (Expr.add (Expr.var 0) (Expr.var 1)) (Expr.add (Expr.var 1) (Expr.var 2)))
     (Expr.add (Expr.var 2) (Expr.var 1))
     (Expr.add (Expr.var 0) (Expr.var 1))
@@ -31,18 +43,18 @@ example (x₁ x₂ x₃ : Nat) : ((x₁ + x₂) + (x₂ + x₃) < x₃ + x₂) =
     rfl
 
 example (x₁ x₂ : Nat) : x₁ + 2 ≤ 3*x₂ → 4*x₂ ≤ 3 + x₁ → 3 ≤ 2*x₂ → False :=
-  Certificate.of_combine_isUnsat [x₁, x₂]
+  Certificate.of_combine_isUnsat #R[x₁, x₂]
     [ (1, { eq := false, lhs := Expr.add (Expr.var 0) (Expr.num 2), rhs := Expr.mulL 3 (Expr.var 1) }),
       (1, { eq := false, lhs := Expr.mulL 4 (Expr.var 1), rhs := Expr.add (Expr.num 3) (Expr.var 0) }),
       (0, { eq := false, lhs := Expr.num 3, rhs := Expr.mulL 2 (Expr.var 1) }) ]
     rfl
 
 example (x : Nat) (xs : List Nat) : (sizeOf x < 1 + (1 + sizeOf x + sizeOf xs)) = True :=
-  ExprCnstr.eq_true_of_isValid [sizeOf x, sizeOf xs]
+  ExprCnstr.eq_true_of_isValid #R[sizeOf x, sizeOf xs]
     { eq := false, lhs := Expr.inc (Expr.var 0), rhs := Expr.add (Expr.num 1) (Expr.add (Expr.add (Expr.num 1) (Expr.var 0)) (Expr.var 1)) }
     rfl
 
 example (x : Nat) (xs : List Nat) : (1 + (1 + sizeOf x + sizeOf xs) < sizeOf x) = False :=
-  ExprCnstr.eq_false_of_isUnsat [sizeOf x, sizeOf xs]
+  ExprCnstr.eq_false_of_isUnsat #R[sizeOf x, sizeOf xs]
     { eq := false, lhs := Expr.inc <| Expr.add (Expr.num 1) (Expr.add (Expr.add (Expr.num 1) (Expr.var 0)) (Expr.var 1)), rhs := Expr.var 0 }
     rfl

tests/lean/run/som1.lean

@@ -1,6 +1,18 @@
+import Lean
+
 open Nat.SOM
+
+-- Convenient RArray literals
+elab tk:"#R[" ts:term,* "]" : term => do
+  let ts : Array Lean.Syntax := ts
+  let es ← ts.mapM fun stx => Lean.Elab.Term.elabTerm stx none
+  if h : 0 < es.size then
+    return (Lean.RArray.toExpr (← Lean.Meta.inferType es[0]!) id (Lean.RArray.ofArray es h))
+  else
+    throwErrorAt tk "RArray cannot be empty"
+
 example : (x + y) * (x + y + 1) = x * (1 + y + x) + (y + 1 + x) * y :=
-  let ctx := [x, y]
+  let ctx := #R[x, y]
   let lhs : Expr := .mul (.add (.var 0) (.var 1)) (.add (.add (.var 0) (.var 1)) (.num 1))
   let rhs : Expr := .add (.mul (.var 0) (.add (.add (.num 1) (.var 1)) (.var 0)))
                          (.mul (.add (.add (.var 1) (.num 1)) (.var 0)) (.var 1))