Lean4
/-- `linearFormOfExpr red map e` computes the linear form of `e`.
`map` is a lookup map from atomic expressions to variable numbers.
If a new atomic expression is encountered, it is added to the map with a new number.
It matches atomic expressions up to reducibility given by `red`.
Because it matches up to definitional equality, this function must be in the `MetaM` monad,
and forces some functions that call it into `MetaM` as well.
-/
partial def linearFormOfExpr (red : TransparencyMode) (m : ExprMap) (e : Expr) : MetaM (ExprMap × Sum) := do
let e ← whnfR e
match e.numeral? with
| some 0 =>
return ⟨m, TreeMap.empty⟩
| some (n + 1) =>
return ⟨m, scalar (n + 1)⟩
| none =>
match e.getAppFnArgs with
| (``HMul.hMul, #[_, _, _, _, e1, e2]) =>
do
let (m1, comp1) ← linearFormOfExpr red m e1
let (m2, comp2) ← linearFormOfExpr red m1 e2
return (m2, comp1.mul comp2)
| (``HAdd.hAdd, #[_, _, _, _, e1, e2]) =>
do
let (m1, comp1) ← linearFormOfExpr red m e1
let (m2, comp2) ← linearFormOfExpr red m1 e2
return (m2, comp1 + comp2)
| (``HSub.hSub, #[_, _, _, _, e1, e2]) =>
do
let (m1, comp1) ← linearFormOfExpr red m e1
let (m2, comp2) ← linearFormOfExpr red m1 e2
return (m2, comp1 + comp2.map (fun _ v => -v))
| (``Neg.neg, #[_, _, e]) =>
do
let (m1, comp) ← linearFormOfExpr red m e
return (m1, comp.map (fun _ v => -v))
| (``HPow.hPow, #[_, _, _, _, a, n]) =>
do
match n.numeral? with
| some n =>
do
let (m1, comp) ← linearFormOfExpr red m a
return (m1, comp.pow n)
| none =>
linearFormOfAtom red m e
| _ =>
linearFormOfAtom red m e
