Application: Advanced Expression Analysis🔗

In metaprogramming, you might want to perform a multi-layered analysis of an expression. Below is an example that:
Counts how many times each constant appears.
Records the minimum depth at which each constant first appeared .
Demonstrates merging two frequency maps.
Shows how to filter a map to keep only specific entries.
structure ConstInfo where
  count : Nat
  firstDepth : Nat
deriving Repr, Inhabited

instance : ToString ConstInfo where
  toString info := s!"Count: {info.count}, 
    First Depth: {info.firstDepth}"

/-- 
  An analyzer that tracks both count 
  and the depth of the first occurrence.
  Using MetaM allows us to use 'logInfo' for debugging.
-/
def analyzeExpr (e : Expr) (depth : Nat := 0) 
    (emap : HashMap Name ConstInfo := {}) :
    MetaM (HashMap Name ConstInfo) := do
  match e with
  -- for constants
  | .const name _ =>
      -- Simple debug print to the Infoview
      logInfo m!"Found constant: {name} at depth {depth}"
      return emap.alter name fun
        | some info => 
            some { info with count := info.count + 1 }
        | none      => 
            some { count := 1, firstDepth := depth }
  -- for applications like `f a`
  | .app f a => 
      let emap' ← analyzeExpr f depth emap
      analyzeExpr a depth emap'
  -- for lambda and forall
  | .lam _ _ b _ | .forallE _ _ b _ => 
      analyzeExpr b (depth + 1) emap
  | _ => return emap

/-- 
  We merge two HashMaps.
  If a constant exists in both, we sum the counts 
  and take the minimum depth.
-/
def mergeAnalysis (m1 m2 : HashMap Name ConstInfo) 
    : HashMap Name ConstInfo :=
  m2.fold (init := m1) fun acc name info2 =>
    acc.alter name fun
    | some info1 => some { 
        count := info1.count + info2.count, 
        firstDepth := min info1.firstDepth info2.firstDepth 
      }
    | none => some info2

def egExprCounter : MetaM Unit := do
  -- Using 'toExpr' elaborates numbers
  -- into 'OfNat.ofNat' calls
  let e1 ← mkEq (toExpr 1) (toExpr 1)
  let e2 ← mkEq (toExpr 2) (toExpr 3)

  let analysis1 ← analyzeExpr e1
  logInfo s!"Analysis of e1: 
    {analysis1.toList.map (fun (k, v) => (k, s!"{v}"))}"
  let analysis2 ← analyzeExpr e2
  logInfo s!"Analysis of e2: 
    {analysis2.toList.map (fun (k, v) => (k, s!"{v}"))}"

  -- Merge the results of two different expressions
  let combined := mergeAnalysis analysis1 analysis2
  logInfo s!"Analysis of combined: 
    {combined.toList.map (fun (k, v) => (k, s!"{v}"))}"

  -- filter the map to only include constants that 
  -- appeared more than once
  let frequent := combined.toList.filter (·.2.count > 1) 
    |> HashMap.ofList

  -- Map the values to just the counts for final display
  let finalCounts := frequent.toList.map 
    fun (k, v) => (k, v.count)

  IO.println s!"Frequent constants: {finalCounts}"

Frequent constants: [(Nat, 6), (instOfNatNat, 4), (Eq, 2), (OfNat.ofNat, 4)]
Found constant: Eq at depth 0Found constant: Nat at depth 0Found constant: OfNat.ofNat at depth 0Found constant: Nat at depth 0Found constant: instOfNatNat at depth 0Found constant: OfNat.ofNat at depth 0Found constant: Nat at depth 0Found constant: instOfNatNat at depth 0Analysis of e1: 
    [(Nat, Count: 3, 
    First Depth: 0), (instOfNatNat, Count: 2, 
    First Depth: 0), (Eq, Count: 1, 
    First Depth: 0), (OfNat.ofNat, Count: 2, 
    First Depth: 0)]Found constant: Eq at depth 0Found constant: Nat at depth 0Found constant: OfNat.ofNat at depth 0Found constant: Nat at depth 0Found constant: instOfNatNat at depth 0Found constant: OfNat.ofNat at depth 0Found constant: Nat at depth 0Found constant: instOfNatNat at depth 0Analysis of e2: 
    [(Nat, Count: 3, 
    First Depth: 0), (instOfNatNat, Count: 2, 
    First Depth: 0), (Eq, Count: 1, 
    First Depth: 0), (OfNat.ofNat, Count: 2, 
    First Depth: 0)]Analysis of combined: 
    [(Nat, Count: 6, 
    First Depth: 0), (instOfNatNat, Count: 4, 
    First Depth: 0), (Eq, Count: 2, 
    First Depth: 0), (OfNat.ofNat, Count: 4, 
    First Depth: 0)]#eval egExprCounter