Eqns.lean

/-
Copyright (c) 2021 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
import Lean.Meta.Basic
import Lean.Meta.AppBuilder

namespace Lean.Meta

def GetEqnsFn := Name → MetaM (Option (Array Name))

private builtin_initialize getEqnsFnsRef : IO.Ref (List GetEqnsFn) ← IO.mkRef []

/--
  Register a new function for retrieving equation theorems.
  We generate equations theorems on demand, and they are generated by more than one module.
  For example, the structural and well-founded recursion modules generate them.
  Most recent getters are tried first.

  A getter returns an `Option (Array Name)`. The result is `none` if the getter failed.
  Otherwise, it is a sequence of theorem names where each one of them corresponds to
  an alternative. Example: the definition

  ```
  def f (xs : List Nat) : List Nat :=
    match xs with
    | [] => []
    | x::xs => (x+1)::f xs
  ```
  should have two equational theorems associated with it
  ```
  f [] = []
  ```
  and
  ```
  (x : Nat) → (xs : List Nat) → f (x :: xs) = (x+1) :: f xs
  ```
-/
def registerGetEqnsFn (f : GetEqnsFn) : IO Unit := do
  unless (← initializing) do
    throw (IO.userError "failed to register equation getter, this kind of extension can only be registered during initialization")
  getEqnsFnsRef.modify (f :: ·)

/-- Return true iff `declName` is a definition and its type is not a proposition. -/
private def shouldGenerateEqnThms (declName : Name) : MetaM Bool := do
  if let some (.defnInfo info) := (← getEnv).find? declName then
    return !(← isProp info.type)
  else
    return false

structure EqnsExtState where
  map : PHashMap Name (Array Name) := {}
  deriving Inhabited

/- We generate the equations on demand, and do not save them on .olean files. -/
builtin_initialize eqnsExt : EnvExtension EqnsExtState ←
  registerEnvExtension (pure {})

/--
  Simple equation theorem for nonrecursive definitions.
-/
private def mkSimpleEqThm (declName : Name) : MetaM (Option Name) := do
  if let some (.defnInfo info) := (← getEnv).find? declName then
    lambdaTelescope info.value fun xs body => do
      let lhs := mkAppN (mkConst info.name <| info.levelParams.map mkLevelParam) xs
      let type  ← mkForallFVars xs (← mkEq lhs body)
      let value ← mkLambdaFVars xs (← mkEqRefl lhs)
      let name := mkPrivateName (← getEnv) declName ++ `_eq_1
      addDecl <| Declaration.thmDecl {
        name, type, value
        levelParams := info.levelParams
      }
      return some name
  else
    return none

/--
  Return equation theorems for the given declaration.
  By default, we not create equation theorems for nonrecursive definitions.
  You can use `nonRec := true` to override this behavior, a dummy `rfl` proof is created on the fly.
-/
def getEqnsFor? (declName : Name) (nonRec := false) : MetaM (Option (Array Name)) := withLCtx {} {} do
  if let some eqs := eqnsExt.getState (← getEnv) |>.map.find? declName then
    return some eqs
  else if (← shouldGenerateEqnThms declName) then
    for f in (← getEqnsFnsRef.get) do
      if let some r ← f declName then
        modifyEnv fun env => eqnsExt.modifyState env fun s => { s with map := s.map.insert declName r }
        return some r
    if nonRec then
      let some eqThm ← mkSimpleEqThm declName | return none
      let r := #[eqThm]
      modifyEnv fun env => eqnsExt.modifyState env fun s => { s with map := s.map.insert declName r }
      return some r
  return none

def GetUnfoldEqnFn := Name → MetaM (Option Name)

private builtin_initialize getUnfoldEqnFnsRef : IO.Ref (List GetUnfoldEqnFn) ← IO.mkRef []

/--
  Register a new function for retrieving a "unfold" equation theorem.

  We generate this kind of equation theorem on demand, and it is generated by more than one module.
  For example, the structural and well-founded recursion modules generate it.
  Most recent getters are tried first.

  A getter returns an `Option Name`. The result is `none` if the getter failed.
  Otherwise, it is a theorem name. Example: the definition

  ```
  def f (xs : List Nat) : List Nat :=
    match xs with
    | [] => []
    | x::xs => (x+1)::f xs
  ```
  should have the theorem
  ```
  (xs : Nat) →
    f xs =
      match xs with
      | [] => []
      | x::xs => (x+1)::f xs
  ```
-/
def registerGetUnfoldEqnFn (f : GetUnfoldEqnFn) : IO Unit := do
  unless (← initializing) do
    throw (IO.userError "failed to register equation getter, this kind of extension can only be registered during initialization")
  getUnfoldEqnFnsRef.modify (f :: ·)

/--
  Return a "unfold" theorem for the given declaration.
  By default, we not create unfold theorems for nonrecursive definitions.
  You can use `nonRec := true` to override this behavior.
-/
def getUnfoldEqnFor? (declName : Name) (nonRec := false) : MetaM (Option Name) := withLCtx {} {} do
  if (← shouldGenerateEqnThms declName) then
    for f in (← getUnfoldEqnFnsRef.get) do
      if let some r ← f declName then
        return some r
    if nonRec then
      let some #[eqThm] ← getEqnsFor? declName (nonRec := true) | return none
      return some eqThm
   return none

end Lean.Meta