path: root/backends/lean/Base/Diverge/Elab.lean

import Lean
import Lean.Meta.Tactic.Simp
import Init.Data.List.Basic
import Mathlib.Tactic.RunCmd
import Base.Diverge.Base
import Base.Diverge.ElabBase

namespace Diverge

/- Automating the generation of the encoding and the proofs so as to use nice
   syntactic sugar. -/

syntax (name := divergentDef)
  declModifiers "divergent" "def" declId ppIndent(optDeclSig) declVal : command

open Lean Elab Term Meta Primitives

set_option trace.Diverge.def true

/- The following was copied from the `wfRecursion` function. -/

open WF in



-- Replace the recursive calls by a call to the continuation
-- def replace_rec_calls

#check Lean.Meta.forallTelescope
#check Expr
#check withRef
#check MonadRef.withRef
#check Nat
#check Array
#check Lean.Meta.inferType
#check Nat
#check Int

#check (0, 1)
#check Prod
#check ()
#check Unit
#check Sigma

--  print_decl is_even_body
#check instOfNatNat
#check OfNat.ofNat -- @OfNat.ofNat ℕ 2 ...
#check OfNat.ofNat -- @OfNat.ofNat (Fin 2) 1 ...
#check Fin.instOfNatFinHAddNatInstHAddInstAddNatOfNat


-- TODO: is there already such a utility somewhere?
-- TODO: change to mkSigmas
def mkProds (tys : List Expr) : MetaM Expr :=
  match tys with
  | [] => do return (Expr.const ``PUnit.unit [])
  | [ty] => do return ty
  | ty :: tys => do
    let pty ← mkProds tys
    mkAppM ``Prod.mk #[ty, pty]

def divRecursion (preDefs : Array PreDefinition) : TermElabM Unit := do
  let msg := toMessageData <| preDefs.map fun pd => (pd.declName, pd.levelParams, pd.type, pd.value)
  trace[Diverge.def] ("divRecursion: defs: " ++ msg)

  -- CHANGE HERE This function should add definitions with these names/types/values ^^
  -- Temporarily add the predefinitions as axioms
  for preDef in preDefs do
    addAsAxiom preDef

  -- TODO: what is this?
  for preDef in preDefs do
    applyAttributesOf #[preDef] AttributeApplicationTime.afterCompilation

  -- Retrieve the name of the first definition, that we will use as the namespace
  -- for the definitions common to the group
  let def0 := preDefs[0]!
  let grName := def0.declName
  trace[Diverge.def] "group name: {grName}"

  /- Compute the type of the continuation.

     We do the following
     - we make sure all the definitions have the same universe parameters
       (we can make this more general later)
     - we group all the type parameters together, make sure all the
       definitions have the same type parameters, and enforce
       a uniform polymorphism (we can also lift this later).
       This would require generalizing a bit our indexed fixed point to
       make the output type parametric in the input.
     - we group all the non-type parameters: we parameterize the continuation
       by those
  -/
  let grLvlParams := def0.levelParams
  trace[Diverge.def] "def0 type: {def0.type}"

  -- Small utility: compute the list of type parameters
  let getTypeParams (ty: Expr) : MetaM (List Expr × List Expr × Expr) :=
    Lean.Meta.forallTelescope ty fun tys out_ty => do
      trace[Diverge.def] "types: {tys}"
/-      let (_, params) ← StateT.run (do
        for x in tys do
          let ty ← Lean.Meta.inferType x
          match ty with
          | .sort _ => do
            let st ← StateT.get
            StateT.set (ty :: st)
          | _ => do break
        ) ([] : List Expr)
      let params := params.reverse
      trace[Diverge.def] "  type parameters {params}"
      return params -/
      let rec get_params (ls : List Expr) : MetaM (List Expr × List Expr) :=
        match ls with
        | x :: tl => do
          let ty ← Lean.Meta.inferType x
          match ty with
          | .sort _ => do
            let (ty_params, params) ← get_params tl
            return (x :: ty_params, params)
          | _ => do return ([], ls)
        | _ => do return ([], [])
      let (ty_params, params) ← get_params tys.toList
      trace[Diverge.def] "  parameters: {ty_params}; {params}"
      return (ty_params, params, out_ty)
  let (grTyParams, _, _) ← do
    getTypeParams def0.type

  -- Compute the input types and the output types
  let all_tys ← preDefs.mapM fun preDef => do
    let (tyParams, params, ret_ty) ← getTypeParams preDef.type
    -- TODO: this is not complete, there are more checks to perform
    if tyParams.length ≠ grTyParams.length then
      throwError "Non-uniform polymorphism"
    return (params, ret_ty)

  -- TODO: I think there are issues with the free variables
  let (input_tys, output_tys) := List.unzip all_tys.toList
  let input_tys : List Expr ← liftM (List.mapM mkProds input_tys)

  trace[Diverge.def] "  in/out tys: {input_tys}; {output_tys}"

  -- Compute the names set
  let names := preDefs.map PreDefinition.declName
  let names := HashSet.empty.insertMany names

  --
  for preDef in preDefs do
    trace[Diverge.def] "about to explore: {preDef.declName}"
    explore_term "" preDef.value

  -- Compute the bodies

  -- Process the definitions
  addAndCompilePartialRec preDefs

-- The following function is copy&pasted from Lean.Elab.PreDefinition.Main
-- This is the only part where we make actual changes and hook into the equation compiler.
-- (I've removed all the well-founded stuff to make it easier to read though.)

open private ensureNoUnassignedMVarsAtPreDef betaReduceLetRecApps partitionPreDefs
  addAndCompilePartial addAsAxioms from Lean.Elab.PreDefinition.Main

def addPreDefinitions (preDefs : Array PreDefinition) : TermElabM Unit := withLCtx {} {} do
  for preDef in preDefs do
    trace[Diverge.elab] "{preDef.declName} : {preDef.type} :=\n{preDef.value}"
  let preDefs ← preDefs.mapM ensureNoUnassignedMVarsAtPreDef
  let preDefs ← betaReduceLetRecApps preDefs
  let cliques := partitionPreDefs preDefs
  let mut hasErrors := false
  for preDefs in cliques do
    trace[Diverge.elab] "{preDefs.map (·.declName)}"
    try
      trace[Diverge.elab] "calling divRecursion"
      withRef (preDefs[0]!.ref) do
        divRecursion preDefs
      trace[Diverge.elab] "divRecursion succeeded"
    catch ex =>
      -- If it failed, we 
      trace[Diverge.elab] "divRecursion failed"
      hasErrors := true
      logException ex
      let s ← saveState
      try
        if preDefs.all fun preDef => preDef.kind == DefKind.def ||
           preDefs.all fun preDef => preDef.kind == DefKind.abbrev then
          -- try to add as partial definition
          try
            addAndCompilePartial preDefs (useSorry := true)
          catch _ =>
            -- Compilation failed try again just as axiom
            s.restore
            addAsAxioms preDefs
        else return ()
      catch _ => s.restore

-- The following two functions are copy&pasted from Lean.Elab.MutualDef

open private elabHeaders levelMVarToParamHeaders getAllUserLevelNames withFunLocalDecls elabFunValues
  instantiateMVarsAtHeader instantiateMVarsAtLetRecToLift checkLetRecsToLiftTypes withUsed from Lean.Elab.MutualDef

def Term.elabMutualDef (vars : Array Expr) (views : Array DefView) : TermElabM Unit := do
    let scopeLevelNames ← getLevelNames
    let headers ← elabHeaders views
    let headers ← levelMVarToParamHeaders views headers
    let allUserLevelNames := getAllUserLevelNames headers
    withFunLocalDecls headers fun funFVars => do
      for view in views, funFVar in funFVars do
        addLocalVarInfo view.declId funFVar
      let values ←
        try
          let values ← elabFunValues headers
          Term.synthesizeSyntheticMVarsNoPostponing
          values.mapM (instantiateMVars ·)
        catch ex =>
          logException ex
          headers.mapM fun header => mkSorry header.type (synthetic := true)
      let headers ← headers.mapM instantiateMVarsAtHeader
      let letRecsToLift ← getLetRecsToLift
      let letRecsToLift ← letRecsToLift.mapM instantiateMVarsAtLetRecToLift
      checkLetRecsToLiftTypes funFVars letRecsToLift
      withUsed vars headers values letRecsToLift fun vars => do
        let preDefs ← MutualClosure.main vars headers funFVars values letRecsToLift
        for preDef in preDefs do
          trace[Diverge.elab] "{preDef.declName} : {preDef.type} :=\n{preDef.value}"
        let preDefs ← withLevelNames allUserLevelNames <| levelMVarToParamPreDecls preDefs
        let preDefs ← instantiateMVarsAtPreDecls preDefs
        let preDefs ← fixLevelParams preDefs scopeLevelNames allUserLevelNames
        for preDef in preDefs do
          trace[Diverge.elab] "after eraseAuxDiscr, {preDef.declName} : {preDef.type} :=\n{preDef.value}"
        checkForHiddenUnivLevels allUserLevelNames preDefs
        addPreDefinitions preDefs

open Command in
def Command.elabMutualDef (ds : Array Syntax) : CommandElabM Unit := do
  let views ← ds.mapM fun d => do
    let `($mods:declModifiers divergent def $id:declId $sig:optDeclSig $val:declVal) := d
      | throwUnsupportedSyntax
    let modifiers ← elabModifiers mods
    let (binders, type) := expandOptDeclSig sig
    let deriving? := none
    pure { ref := d, kind := DefKind.def, modifiers,
           declId := id, binders, type? := type, value := val, deriving? }
  runTermElabM fun vars => Term.elabMutualDef vars views

-- Special command so that we don't fall back to the built-in mutual when we produce an error.
local syntax "_divergent" Parser.Command.mutual : command
elab_rules : command | `(_divergent mutual $decls* end) => Command.elabMutualDef decls

macro_rules
  | `(mutual $decls* end) => do
    unless !decls.isEmpty && decls.all (·.1.getKind == ``divergentDef) do
      Macro.throwUnsupported
    `(command| _divergent mutual $decls* end)

open private setDeclIdName from Lean.Elab.Declaration
elab_rules : command
  | `($mods:declModifiers divergent%$tk def $id:declId $sig:optDeclSig $val:declVal) => do
    let (name, _) := expandDeclIdCore id
    if (`_root_).isPrefixOf name then throwUnsupportedSyntax
    let view := extractMacroScopes name
    let .str ns shortName := view.name | throwUnsupportedSyntax
    let shortName' := { view with name := shortName }.review
    let cmd ← `(mutual $mods:declModifiers divergent%$tk def $(⟨setDeclIdName id shortName'⟩):declId $sig:optDeclSig $val:declVal end)
    if ns matches .anonymous then
      Command.elabCommand cmd
    else
      Command.elabCommand <| ← `(namespace $(mkIdentFrom id ns) $cmd end $(mkIdentFrom id ns))

mutual
  divergent def is_even (i : Int) : Result Bool :=
    if i = 0 then return true else return (← is_odd (i - 1))

  divergent def is_odd (i : Int) : Result Bool :=
    if i = 0 then return false else return (← is_even (i - 1))
end

example (i : Int) : is_even i = .ret (i % 2 = 0) ∧ is_odd i = .ret (i % 2 ≠ 0) := by
  induction i
  unfold is_even
  sorry    

divergent def list_nth {a: Type} (ls : List a) (i : Int) : Result a :=
  match ls with
  | [] => .fail .panic
  | x :: ls =>
    if i = 0 then return x
    else return (← list_nth ls (i - 1))

mutual
  divergent def foo (i : Int) : Result Nat :=
    if i > 10 then return (← foo (i / 10)) + (← bar i) else bar 10

  divergent def bar (i : Int) : Result Nat :=
    if i > 20 then foo (i / 20) else .ret 42
end

end Diverge