import Lean
import Lean.Parser.Command
import Lean.Parser.Term

namespace NamedState

open Lean Elab Command Term Meta TSyntax

--
-- A monad to read/write from multiples state variables, using string lookup.
-- As an example if your state is the struct { z : Nat, y : String } then you
-- can (getNamed "y") or (putNamed "z" 3) in your do block.
-- If using algebraic effects you would need something like [Sendable (NamedState "z" Nat) m]
-- to getNamed/putNamed z.
--

inductive NamedState (n : String) (v : Type) : Type → Type where
  | Get : NamedState n v v
  | Put : v → NamedState n v Unit

/-
 A general monad to use as the target for a collapser.
 Usually when running algebraic effects, the monads you end up with are IO and
 some state. So you can build your collapser to produce a StateIO which handles both of these cases.

 Lets' say you have a sum type and freer monad or some equivalent construct. Here we make a freer monad with
 two effects, (NamedState "z" Nat α) and (IO α)

 > mkSumType ExampleCommand >| (NamedState "z" Nat), IO |<
 > mkFreer ExampleMonad ExampleCommand

 This has two sendable instances, IO and (NamedState "z" Nat).

 To interpret this we make a StateIO monad as the final target monad:

 > mkStateIO Blargh (z:Nat),(y:String) @@

 this makes a datatype of type "StateIO Blarghstruct" where Blarghstruct is a struct with fields { z:Nat, y:String }

 Then in the interpreter you can use "collapseNamedState" for example:

 > def interpreter1 := buildInterpreter ExampleCommand OneState (NamedState "z" Nat),IO
 >   [:
 >     collapseNamedState "z" Nat,
 >     collapseIO
 >   :]

-/

def StateIO (sType : Type) (α : Type) : Type := sType → IO (α × sType)

instance : Monad (StateIO s) where
    pure := fun a s => pure ⟨a, s⟩
    bind := fun m f s => do let ⟨a', s'⟩ ← m s
                            f a' s'

class StateOperator (stateContainer : Type) (name : String) (state : Type) where
    putS : state → stateContainer → stateContainer
    getS : stateContainer → state

-- Normally a state monad has a single variable that you access using get/put.
-- This builds a structure representing state, with several fields in it. Each field "x"
-- is a single state that is accessed using putNamed "x"/getNamed "x".
set_option hygiene false in
def elabSS (structid : TSyntax `Lean.Parser.Command.declId) (vals : Syntax.TSepArray `structfield ",") : CommandElabM Unit := do
    let valArray : Array (TSyntax `structfield) := vals
    let valInstance : TSyntax `structfield → CommandElabM (TSyntax `Lean.Parser.Command.structExplicitBinder) :=
      fun n => do
        let id : Ident := TSyntax.mk <| n.raw.getArgs[1]!
        let ftype : Term := TSyntax.mk <| n.raw.getArgs[3]!
        let c ← `(Lean.Parser.Command.structExplicitBinder | ($id : $ftype))
        pure c
    let fields ← Array.sequenceMap valArray valInstance
    let structDecl ← `(structure $structid where $fields:structExplicitBinder*)
    elabCommand structDecl

declare_syntax_cat structfield
syntax " ( " ident " : " term " ) " : structfield

elab "mkStateIOStruct" structid:ident vals:structfield,+ " @@ " : command => elabSS structid vals

-- This makes instances of StateOperator for a particular state container (a structure) and a named
-- field of that structure. There should be an instance generated for each field of the structure.
set_option hygiene false in
def elabSI (structid :Term) (fields : Syntax.TSepArray `structfield ",") : CommandElabM Unit := do
  let fieldArray : Array (TSyntax `structfield) := fields
  let fieldInstance : TSyntax `structfield → CommandElabM Unit :=
    fun n => do
      let id := TSyntax.mk n.raw.getArgs[1]!
      let ftype : Term := TSyntax.mk n.raw.getArgs[3]!
      let s := Syntax.mkStrLit id.getId.toString
      let c ← `(instance : StateOperator $structid $s $ftype where
                  putS := fun v s => { s with $id:ident := v}
                  getS := fun s => s.$id)
      elabCommand c
  Array.forM fieldInstance fieldArray

elab "mkStateInterfaces" structid:term vals:structfield,+ " @@ " : command => elabSI structid vals

-- Makes a complete set of definitions for a StateIO monad, including:
--  A structure with fields to hold all the named states
--  Instances of StateOperator to get/put state
--  A Monad instance
-- You provide the monad name and field names/types.  For a StateIO monad named "x" there is also
-- a State structure named "xstruct" which you can use.
elab "mkStateIO" stateIOname:ident vals:structfield,+ " @@ " : command => do
    let structid : Ident := Lean.mkIdent <| Name.appendAfter stateIOname.getId "struct"
    elabSS structid vals
    elabSI structid vals
    let siodef ← `(def $stateIOname := StateIO $structid)
    elabCommand siodef
    let c ←
        `(instance : Monad $stateIOname where
              pure := fun a s => pure ⟨a, s⟩
              bind := fun m f s => do let ⟨a', s'⟩ ← m s
                                      f a' s')
    elabCommand c

--mkStateIOStruct Blargh (z:Nat),(y:String) @@
--mkStateInterfaces Blargh (z:Nat),(y:String) @@

/-
mkStateIO Blargh (z:Nat),(y:String) @@

def testStruct : Blarghstruct := { z := 3, y := "argh"}

#check testStruct
def goP [StateOperator Blarghstruct "z" Nat] : Blarghstruct → Nat := fun b => StateOperator.getS "z" b
#eval goP testStruct
-/

-- When running the interpreter for a Freer monad that has one or more named state, you typically collapse
-- the monad into a StateIO monad.  For each effect that is a NamedState you can use collapseNamedState.
def collapseNamedState (n : String) (v : Type) [StateOperator s n v] {α : Type} : NamedState n v α → StateIO s α :=
  fun m =>
    match m with
    | .Get => fun s => pure ⟨StateOperator.getS n s,s⟩
    | .Put v' => fun s => pure ⟨(), StateOperator.putS n v' s⟩

-- When running the interpreter for a Freer monad that has one or more named state, you typically collapse
-- the monad into a StateIO monad.  If you have an effect that represents arbitrary IO you can collapse it with collapseIO.
def collapseIO : IO α → StateIO ss α :=
    fun o => fun s => Functor.map (fun x => ⟨x,s⟩) o

end NamedState


/-
Copyright (c) 2021 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Sebastian Ullrich, Leonardo de Moura
-/
prelude
import Init.SimpLemmas
import Init.Control.Except
import Init.Control.StateRef

open Function

@[simp] theorem monadLift_self [Monad m] (x : m α) : monadLift x = x :=
  rfl

class LawfulFunctor (f : Type u → Type v) [Functor f] : Prop where
  map_const          : (Functor.mapConst : α → f β → f α) = Functor.map ∘ const β
  id_map   (x : f α) : id <$> x = x
  comp_map (g : α → β) (h : β → γ) (x : f α) : (h ∘ g) <$> x = h <$> g <$> x

export LawfulFunctor (map_const id_map comp_map)

attribute [simp] id_map

@[simp] theorem id_map' [Functor m] [LawfulFunctor m] (x : m α) : (fun a => a) <$> x = x :=
  id_map x

class LawfulApplicative (f : Type u → Type v) [Applicative f] extends LawfulFunctor f : Prop where
  seqLeft_eq  (x : f α) (y : f β)     : x <* y = const β <$> x <*> y
  seqRight_eq (x : f α) (y : f β)     : x *> y = const α id <$> x <*> y
  pure_seq    (g : α → β) (x : f α)   : pure g <*> x = g <$> x
  map_pure    (g : α → β) (x : α)     : g <$> (pure x : f α) = pure (g x)
  seq_pure    {α β : Type u} (g : f (α → β)) (x : α) : g <*> pure x = (fun h => h x) <$> g
  seq_assoc   {α β γ : Type u} (x : f α) (g : f (α → β)) (h : f (β → γ)) : h <*> (g <*> x) = ((@comp α β γ) <$> h) <*> g <*> x
  comp_map g h x := (by
    repeat rw [← pure_seq]
    simp [seq_assoc, map_pure, seq_pure])

export LawfulApplicative (seqLeft_eq seqRight_eq pure_seq map_pure seq_pure seq_assoc)

attribute [simp] map_pure seq_pure

@[simp] theorem pure_id_seq [Applicative f] [LawfulApplicative f] (x : f α) : pure id <*> x = x := by
  simp [pure_seq]

class LawfulMonad (m : Type u → Type v) [Monad m] extends LawfulApplicative m : Prop where
  bind_pure_comp (f : α → β) (x : m α) : x >>= pure ∘ f = f <$> x
  bind_map       {α β : Type u} (f : m (α → β)) (x : m α) : f >>= (. <$> x) = f <*> x
  pure_bind      (x : α) (f : α → m β) : pure x >>= f = f x
  bind_assoc     (x : m α) (f : α → m β) (g : β → m γ) : x >>= f >>= g = x >>= fun x => f x >>= g
  map_pure g x    := (by rw [← bind_pure_comp, pure_bind])
  seq_pure g x    := (by rw [← bind_map]; simp [map_pure, bind_pure_comp])
  seq_assoc x g h := (by
    -- TODO: support for applying `symm` at `simp` arguments
    let bind_pure_comp_symm {α β : Type u} (f : α → β) (x : m α) : f <$> x = x >>= pure ∘ f := by
      rw [bind_pure_comp]
    let bind_map_symm {α β : Type u} (f : m (α → (β : Type u))) (x : m α) : f <*> x = f >>= (. <$> x) := by
      rw [bind_map]
    simp[bind_pure_comp_symm, bind_map_symm, bind_assoc, pure_bind])

export LawfulMonad (bind_pure_comp bind_map pure_bind bind_assoc)
attribute [simp] pure_bind bind_assoc

@[simp] theorem bind_pure [Monad m] [LawfulMonad m] (x : m α) : x >>= pure = x := by
  show x >>= pure ∘ id = x
  rw [bind_pure_comp, id_map]

theorem map_eq_pure_bind [Monad m] [LawfulMonad m] (f : α → β) (x : m α) : f <$> x = x >>= fun a => pure (f a) := by
  rw [← bind_pure_comp]

theorem seq_eq_bind_map {α β : Type u} [Monad m] [LawfulMonad m] (f : m (α → β)) (x : m α) : f <*> x = f >>= (. <$> x) := by
  rw [← bind_map]

theorem bind_congr [Bind m] {x : m α} {f g : α → m β} (h : ∀ a, f a = g a) : x >>= f = x >>= g := by
  simp [funext h]

@[simp] theorem bind_pure_unit [Monad m] [LawfulMonad m] {x : m PUnit} : (x >>= fun _ => pure ⟨⟩) = x := by
  have : (x >>= fun _ => pure ⟨⟩) = (x >>= pure) := by
    apply bind_congr; intro u
    cases u; simp
  rw [bind_pure] at this
  assumption

theorem map_congr [Functor m] {x : m α} {f g : α → β} (h : ∀ a, f a = g a) : (f <$> x : m β) = g <$> x := by
  simp [funext h]

theorem seq_eq_bind {α β : Type u} [Monad m] [LawfulMonad m] (mf : m (α → β)) (x : m α) : mf <*> x = mf >>= fun f => f <$> x := by
  rw [bind_map]

theorem seqRight_eq_bind [Monad m] [LawfulMonad m] (x : m α) (y : m β) : x *> y = x >>= fun _ => y := by
  rw [seqRight_eq]
  simp [map_eq_pure_bind, seq_eq_bind_map, const]

theorem seqLeft_eq_bind [Monad m] [LawfulMonad m] (x : m α) (y : m β) : x <* y = x >>= fun a => y >>= fun _ => pure a := by
  rw [seqLeft_eq]; simp [map_eq_pure_bind, seq_eq_bind_map]

/- Id -/

namespace Id

@[simp] theorem map_eq (x : Id α) (f : α → β) : f <$> x = f x := rfl
@[simp] theorem bind_eq (x : Id α) (f : α → id β) : x >>= f = f x := rfl
@[simp] theorem pure_eq (a : α) : (pure a : Id α) = a := rfl

instance : LawfulMonad Id := by
  refine' { .. } <;> intros <;> rfl

end Id

/- ExceptT -/

namespace ExceptT

theorem ext [Monad m] {x y : ExceptT ε m α} (h : x.run = y.run) : x = y := by
  simp [run] at h
  assumption

@[simp] theorem run_pure [Monad m] : run (pure x : ExceptT ε m α) = pure (Except.ok x) := rfl

@[simp] theorem run_lift [Monad m] (x : m α) : run (ExceptT.lift x : ExceptT ε m α) = (Except.ok <$> x : m (Except ε α)) := rfl

@[simp] theorem run_throw [Monad m] : run (throw e : ExceptT ε m β) = pure (Except.error e) := rfl

@[simp] theorem run_bind_lift [Monad m] [LawfulMonad m] (x : m α) (f : α → ExceptT ε m β) : run (ExceptT.lift x >>= f : ExceptT ε m β) = x >>= fun a => run (f a) := by
  simp[ExceptT.run, ExceptT.lift, bind, ExceptT.bind, ExceptT.mk, ExceptT.bindCont, map_eq_pure_bind]

@[simp] theorem bind_throw [Monad m] [LawfulMonad m] (f : α → ExceptT ε m β) : (throw e >>= f) = throw e := by
  simp [throw, throwThe, MonadExceptOf.throw, bind, ExceptT.bind, ExceptT.bindCont, ExceptT.mk]

theorem run_bind [Monad m] (x : ExceptT ε m α)
        : run (x >>= f : ExceptT ε m β)
          =
          run x >>= fun
                     | Except.ok x => run (f x)
                     | Except.error e => pure (Except.error e) :=
  rfl

@[simp] theorem lift_pure [Monad m] [LawfulMonad m] (a : α) : ExceptT.lift (pure a) = (pure a : ExceptT ε m α) := by
  simp [ExceptT.lift, pure, ExceptT.pure]

@[simp] theorem run_map [Monad m] [LawfulMonad m] (f : α → β) (x : ExceptT ε m α)
    : (f <$> x).run = Except.map f <$> x.run := by
  simp [Functor.map, ExceptT.map, map_eq_pure_bind]
  apply bind_congr
  intro a; cases a <;> simp [Except.map]

protected theorem seq_eq {α β ε : Type u} [Monad m] (mf : ExceptT ε m (α → β)) (x : ExceptT ε m α) : mf <*> x = mf >>= fun f => f <$> x :=
  rfl

protected theorem bind_pure_comp [Monad m] [LawfulMonad m] (f : α → β) (x : ExceptT ε m α) : x >>= pure ∘ f = f <$> x := by
  intros; rfl

protected theorem seqLeft_eq {α β ε : Type u} {m : Type u → Type v} [Monad m] [LawfulMonad m] (x : ExceptT ε m α) (y : ExceptT ε m β) : x <* y = const β <$> x <*> y := by
  show (x >>= fun a => y >>= fun _ => pure a) = (const (α := α) β <$> x) >>= fun f => f <$> y
  rw [← ExceptT.bind_pure_comp]
  apply ext
  simp [run_bind]
  apply bind_congr
  intro
  | Except.error _ => simp
  | Except.ok _ =>
    simp [map_eq_pure_bind]; apply bind_congr; intro b;
    cases b <;> simp [comp, Except.map, const]

protected theorem seqRight_eq [Monad m] [LawfulMonad m] (x : ExceptT ε m α) (y : ExceptT ε m β) : x *> y = const α id <$> x <*> y := by
  show (x >>= fun _ => y) = (const α id <$> x) >>= fun f => f <$> y
  rw [← ExceptT.bind_pure_comp]
  apply ext
  simp [run_bind]
  apply bind_congr
  intro a; cases a <;> simp

instance [Monad m] [LawfulMonad m] : LawfulMonad (ExceptT ε m) where
  id_map         := by intros; apply ext; simp
  map_const      := by intros; rfl
  seqLeft_eq     := ExceptT.seqLeft_eq
  seqRight_eq    := ExceptT.seqRight_eq
  pure_seq       := by intros; apply ext; simp [ExceptT.seq_eq, run_bind]
  bind_pure_comp := ExceptT.bind_pure_comp
  bind_map       := by intros; rfl
  pure_bind      := by intros; apply ext; simp [run_bind]
  bind_assoc     := by intros; apply ext; simp [run_bind]; apply bind_congr; intro a; cases a <;> simp

end ExceptT

/- ReaderT -/

namespace ReaderT

theorem ext [Monad m] {x y : ReaderT ρ m α} (h : ∀ ctx, x.run ctx = y.run ctx) : x = y := by
  simp [run] at h
  exact funext h

@[simp] theorem run_pure [Monad m] (a : α) (ctx : ρ) : (pure a : ReaderT ρ m α).run ctx = pure a := rfl

@[simp] theorem run_bind [Monad m] (x : ReaderT ρ m α) (f : α → ReaderT ρ m β) (ctx : ρ)
    : (x >>= f).run ctx = x.run ctx >>= λ a => (f a).run ctx := rfl

@[simp] theorem run_map [Monad m] (f : α → β) (x : ReaderT ρ m α) (ctx : ρ)
    : (f <$> x).run ctx = f <$> x.run ctx := rfl

@[simp] theorem run_monadLift [MonadLiftT n m] (x : n α) (ctx : ρ)
    : (monadLift x : ReaderT ρ m α).run ctx = (monadLift x : m α) := rfl

@[simp] theorem run_monadMap [Monad m] [MonadFunctor n m] (f : {β : Type u} → n β → n β) (x : ReaderT ρ m α) (ctx : ρ)
    : (monadMap @f x : ReaderT ρ m α).run ctx = monadMap @f (x.run ctx) := rfl

@[simp] theorem run_read [Monad m] (ctx : ρ) : (ReaderT.read : ReaderT ρ m ρ).run ctx = pure ctx := rfl

@[simp] theorem run_seq {α β : Type u} [Monad m] [LawfulMonad m] (f : ReaderT ρ m (α → β)) (x : ReaderT ρ m α) (ctx : ρ) : (f <*> x).run ctx = (f.run ctx <*> x.run ctx) := by
  rw [seq_eq_bind (m := m)]; rfl

@[simp] theorem run_seqRight [Monad m] [LawfulMonad m] (x : ReaderT ρ m α) (y : ReaderT ρ m β) (ctx : ρ) : (x *> y).run ctx = (x.run ctx *> y.run ctx) := by
  rw [seqRight_eq_bind (m := m)]; rfl

@[simp] theorem run_seqLeft [Monad m] [LawfulMonad m] (x : ReaderT ρ m α) (y : ReaderT ρ m β) (ctx : ρ) : (x <* y).run ctx = (x.run ctx <* y.run ctx) := by
  rw [seqLeft_eq_bind (m := m)]; rfl

instance [Monad m] [LawfulMonad m] : LawfulMonad (ReaderT ρ m) where
  id_map         := by intros; apply ext; intros; simp
  map_const      := by intros; rfl
  seqLeft_eq     := by intros; apply ext; intros; simp; apply LawfulApplicative.seqLeft_eq
  seqRight_eq    := by intros; apply ext; intros; simp; apply LawfulApplicative.seqRight_eq
  pure_seq       := by intros; apply ext; intros; simp; apply LawfulApplicative.pure_seq
  bind_pure_comp := by intros; apply ext; intros; simp; apply LawfulMonad.bind_pure_comp
  bind_map       := by intros; rfl
  pure_bind      := by intros; apply ext; intros; simp
  bind_assoc     := by intros; apply ext; intros; simp

end ReaderT

/- StateRefT -/

instance [Monad m] [LawfulMonad m] : LawfulMonad (StateRefT' ω σ m) :=
  inferInstanceAs (LawfulMonad (ReaderT (ST.Ref ω σ) m))

/- StateT -/

namespace StateT

theorem ext {x y : StateT σ m α} (h : ∀ s, x.run s = y.run s) : x = y :=
  funext h

@[simp] theorem run'_eq [Monad m] (x : StateT σ m α) (s : σ) : run' x s = (·.1) <$> run x s :=
  rfl

@[simp] theorem run_pure [Monad m] (a : α) (s : σ) : (pure a : StateT σ m α).run s = pure (a, s) := rfl

@[simp] theorem run_bind [Monad m] (x : StateT σ m α) (f : α → StateT σ m β) (s : σ)
    : (x >>= f).run s = x.run s >>= λ p => (f p.1).run p.2 := by
  simp [bind, StateT.bind, run]
  apply bind_congr
  intro p; cases p; rfl

@[simp] theorem run_map {α β σ : Type u} [Monad m] [LawfulMonad m] (f : α → β) (x : StateT σ m α) (s : σ) : (f <$> x).run s = (fun (p : α × σ) => (f p.1, p.2)) <$> x.run s := by
  simp [Functor.map, StateT.map, run, map_eq_pure_bind]
  apply bind_congr
  intro p; cases p; rfl

@[simp] theorem run_get [Monad m] (s : σ)    : (get : StateT σ m σ).run s = pure (s, s) := rfl

@[simp] theorem run_set [Monad m] (s s' : σ) : (set s' : StateT σ m PUnit).run s = pure (⟨⟩, s') := rfl

@[simp] theorem run_modify [Monad m] (f : σ → σ) (s : σ) : (modify f : StateT σ m PUnit).run s = pure (⟨⟩, f s) := rfl

@[simp] theorem run_modifyGet [Monad m] (f : σ → α × σ) (s : σ) : (modifyGet f : StateT σ m α).run s = pure ((f s).1, (f s).2) := by
  simp [modifyGet, MonadStateOf.modifyGet, StateT.modifyGet, run]; cases f s <;> rfl

@[simp] theorem run_lift {α σ : Type u} [Monad m] (x : m α) (s : σ) : (StateT.lift x : StateT σ m α).run s = x >>= fun a => pure (a, s) := rfl

@[simp] theorem run_bind_lift {α σ : Type u} [Monad m] [LawfulMonad m] (x : m α) (f : α → StateT σ m β) (s : σ) : (StateT.lift x >>= f).run s = x >>= fun a => (f a).run s := by
  simp [StateT.lift, StateT.run, bind, StateT.bind]

@[simp] theorem run_monadLift {α σ : Type u} [Monad m] [MonadLiftT n m] (x : n α) (s : σ) : (monadLift x : StateT σ m α).run s = (monadLift x : m α) >>= fun a => pure (a, s) := rfl

@[simp] theorem run_monadMap [Monad m] [MonadFunctor n m] (f : {β : Type u} → n β → n β) (x : StateT σ m α) (s : σ)
    : (monadMap @f x : StateT σ m α).run s = monadMap @f (x.run s) := rfl

@[simp] theorem run_seq {α β σ : Type u} [Monad m] [LawfulMonad m] (f : StateT σ m (α → β)) (x : StateT σ m α) (s : σ) : (f <*> x).run s = (f.run s >>= fun fs => (fun (p : α × σ) => (fs.1 p.1, p.2)) <$> x.run fs.2) := by
  show (f >>= fun g => g <$> x).run s = _
  simp

@[simp] theorem run_seqRight [Monad m] [LawfulMonad m] (x : StateT σ m α) (y : StateT σ m β) (s : σ) : (x *> y).run s = (x.run s >>= fun p => y.run p.2) := by
  show (x >>= fun _ => y).run s = _
  simp

@[simp] theorem run_seqLeft {α β σ : Type u} [Monad m] [LawfulMonad m] (x : StateT σ m α) (y : StateT σ m β) (s : σ) : (x <* y).run s = (x.run s >>= fun p => y.run p.2 >>= fun p' => pure (p.1, p'.2)) := by
  show (x >>= fun a => y >>= fun _ => pure a).run s = _
  simp

theorem seqRight_eq [Monad m] [LawfulMonad m] (x : StateT σ m α) (y : StateT σ m β) : x *> y = const α id <$> x <*> y := by
  apply ext; intro s
  simp [map_eq_pure_bind, const]
  apply bind_congr; intro p; cases p
  simp [Prod.ext]

theorem seqLeft_eq [Monad m] [LawfulMonad m] (x : StateT σ m α) (y : StateT σ m β) : x <* y = const β <$> x <*> y := by
  apply ext; intro s
  simp [map_eq_pure_bind]

instance [Monad m] [LawfulMonad m] : LawfulMonad (StateT σ m) where
  id_map         := by intros; apply ext; intros; simp[Prod.ext]
  map_const      := by intros; rfl
  seqLeft_eq     := seqLeft_eq
  seqRight_eq    := seqRight_eq
  pure_seq       := by intros; apply ext; intros; simp
  bind_pure_comp := by intros; apply ext; intros; simp; apply LawfulMonad.bind_pure_comp
  bind_map       := by intros; rfl
  pure_bind      := by intros; apply ext; intros; simp
  bind_assoc     := by intros; apply ext; intros; simp

end StateT


import init.lean.message init.lean.parser.syntax init.lean.parser.trie init.lean.parser.basic init.lean.parser.stringliteral
import init.lean.parser.token

namespace Lean
namespace flatParser
open String
open Parser (Syntax Syntax.missing)
open Parser (Trie TokenMap)

abbreviation pos := String.Pos

/-- A precomputed cache for quickly mapping Char offsets to positions. -/
structure FileMap :=
(offsets : Array Nat)
(lines   : Array Nat)

namespace FileMap
private def fromStringAux (s : String) : Nat → Nat → Nat → pos → Array Nat → Array Nat → FileMap
| 0     offset line i offsets lines := ⟨offsets.push offset, lines.push line⟩
| (k+1) offset line i offsets lines :=
  if s.atEnd i then ⟨offsets.push offset, lines.push line⟩
  else let c := s.get i in
       let i := s.next i in
       let offset := offset + 1 in
       if c = '\n'
       then fromStringAux k offset (line+1) i (offsets.push offset) (lines.push (line+1))
       else fromStringAux k offset line i offsets lines

def fromString (s : String) : FileMap :=
fromStringAux s s.length 0 1 0 (Array.empty.push 0) (Array.empty.push 1)

/- Remark: `offset is in [(offsets.get b), (offsets.get e)]` and `b < e` -/
private def toPositionAux (offsets : Array Nat) (lines : Array Nat) (offset : Nat) : Nat → Nat → Nat → Position
| 0     b e := ⟨offset, 1⟩ -- unreachable
| (k+1) b e :=
  let offsetB := offsets.get b in
  if e = b + 1 then ⟨offset - offsetB, lines.get b⟩
  else let m := (b + e) / 2 in
       let offsetM := offsets.get m in
       if offset = offsetM then ⟨0, lines.get m⟩
       else if offset > offsetM then toPositionAux k m e
       else toPositionAux k b m

def toPosition : FileMap → Nat → Position
| ⟨offsets, lines⟩ offset := toPositionAux offsets lines offset offsets.size 0 (offsets.size-1)
end FileMap

structure TokenConfig :=
(«prefix» : String)
(lbp : Nat := 0)

structure FrontendConfig :=
(filename : String)
(input    : String)
(FileMap : FileMap)

/- Remark: if we have a Node in the Trie with `some TokenConfig`, the String induced by the path is equal to the `TokenConfig.prefix`. -/
structure ParserConfig extends FrontendConfig :=
(tokens : Trie TokenConfig)

-- Backtrackable State
structure ParserState :=
(messages : MessageLog)

structure TokenCacheEntry :=
(startPos stopPos : pos)
(tk : Syntax)

-- Non-backtrackable State
structure ParserCache :=
(tokenCache : Option TokenCacheEntry := none)

inductive Result (α : Type)
| ok       (a : α)        (i : pos) (cache : ParserCache) (State : ParserState) (eps : Bool) : Result
| error {} (msg : String) (i : pos) (cache : ParserCache) (stx : Syntax)         (eps : Bool) : Result

inductive Result.IsOk {α : Type} : Result α → Prop
| mk (a : α) (i : pos) (cache : ParserCache) (State : ParserState) (eps : Bool) : Result.IsOk (Result.ok a i cache State eps)

theorem errorIsNotOk {α : Type} {msg : String} {i : pos} {cache : ParserCache} {stx : Syntax} {eps : Bool}
                        (h : Result.IsOk (@Result.error α msg i cache stx eps)) : False :=
match h with end

@[inline] def unreachableError {α β : Type} {msg : String} {i : pos} {cache : ParserCache} {stx : Syntax} {eps : Bool}
                                (h : Result.IsOk (@Result.error α msg i cache stx eps)) : β :=
False.elim (errorIsNotOk h)

def resultOk := {r : Result Unit // r.IsOk}

@[inline] def mkResultOk (i : pos) (cache : ParserCache) (State : ParserState) (eps := true) : resultOk :=
⟨Result.ok () i cache State eps, Result.IsOk.mk _ _ _ _ _⟩

def parserCoreM (α : Type) :=
ParserConfig → resultOk → Result α
abbreviation parserCore := parserCoreM Syntax

structure recParsers :=
(cmdParser  : parserCore)
(termParser : Nat → parserCore)

def parserM (α : Type) := recParsers → parserCoreM α
abbreviation Parser := parserM Syntax
abbreviation trailingParser := Syntax → Parser

@[inline] def command.Parser : Parser := λ ps, ps.cmdParser
@[inline] def Term.Parser (rbp : Nat := 0) : Parser  := λ ps, ps.termParser rbp

@[inline] def parserM.pure {α : Type} (a : α) : parserM α :=
λ _ _ r,
  match r with
  | ⟨Result.ok _ it c s _, h⟩   := Result.ok a it c s true
  | ⟨Result.error _ _ _ _ _, h⟩ := unreachableError h

@[inline_if_reduce] def eagerOr  (b₁ b₂ : Bool) := b₁ || b₂
@[inline_if_reduce] def eagerAnd (b₁ b₂ : Bool) := b₁ && b₂

@[inline] def parserM.bind {α β : Type} (x : parserM α) (f : α → parserM β) : parserM β :=
λ ps cfg r,
  match x ps cfg r with
  | Result.ok a i c s e₁ :=
    (match f a ps cfg (mkResultOk i c s) with
     | Result.ok b i c s e₂        := Result.ok b i c s (eagerAnd e₁ e₂)
     | Result.error msg i c stx e₂ := Result.error msg i c stx (eagerAnd e₁ e₂))
  | Result.error msg i c stx e  := Result.error msg i c stx e

instance : Monad parserM :=
{pure := @parserM.pure, bind := @parserM.bind}

@[inline] protected def orelse {α : Type} (p q : parserM α) : parserM α :=
λ ps cfg r,
  match r with
  | ⟨Result.ok _ i₁ _ s₁ _, _⟩ :=
    (match p ps cfg r with
     | Result.error msg₁ i₂ c₂ stx₁ true := q ps cfg (mkResultOk i₁ c₂ s₁)
     | other                           := other)
  | ⟨Result.error _ _ _ _ _, h⟩ := unreachableError h

@[inline] protected def failure {α : Type} : parserM α :=
λ _ _ r,
  match r with
  | ⟨Result.ok _ i c s _, h⟩    := Result.error "failure" i c Syntax.missing true
  | ⟨Result.error _ _ _ _ _, h⟩ := unreachableError h

instance : Alternative parserM :=
{ orelse         := @flatParser.orelse,
  failure        := @flatParser.failure,
  ..flatParser.Monad }

def setSilentError {α : Type} : Result α → Result α
| (Result.error i c msg stx _) := Result.error i c msg stx true
| other                        := other

/--
`try p` behaves like `p`, but it pretends `p` hasn't
consumed any input when `p` fails.
-/
@[inline] def try {α : Type} (p : parserM α) : parserM α :=
λ ps cfg r, setSilentError (p ps cfg r)

@[inline] def atEnd (cfg : ParserConfig) (i : pos) : Bool :=
cfg.input.atEnd i

@[inline] def curr (cfg : ParserConfig) (i : pos) : Char :=
cfg.input.get i

@[inline] def next (cfg : ParserConfig) (i : pos) : pos :=
cfg.input.next i

@[inline] def inputSize (cfg : ParserConfig) : Nat :=
cfg.input.length

@[inline] def currPos : resultOk → pos
| ⟨Result.ok _ i _ _ _, _⟩    := i
| ⟨Result.error _ _ _ _ _, h⟩ := unreachableError h

@[inline] def currState : resultOk → ParserState
| ⟨Result.ok _ _ _ s _, _⟩    := s
| ⟨Result.error _ _ _ _ _, h⟩ := unreachableError h

def mkError {α : Type} (r : resultOk) (msg : String) (stx : Syntax := Syntax.missing) (eps := true) : Result α :=
match r with
| ⟨Result.ok _ i c s _, _⟩    := Result.error msg i c stx eps
| ⟨Result.error _ _ _ _ _, h⟩ := unreachableError h

@[inline] def satisfy (p : Char → Bool) : parserM Char :=
λ _ cfg r,
  match r with
  | ⟨Result.ok _ i ch st e, _⟩ :=
    if atEnd cfg i then mkError r "end of input"
    else let c := curr cfg i in
         if p c then Result.ok c (next cfg i) ch st false
         else mkError r "unexpected character"
  | ⟨Result.error _ _ _ _ _, h⟩ := unreachableError h

def any : parserM Char :=
satisfy (λ _, true)

@[specialize] def takeUntilAux (p : Char → Bool) (cfg : ParserConfig) : Nat → resultOk → Result Unit
| 0     r := r.val
| (n+1) r :=
  match r with
  | ⟨Result.ok _ i ch st e, _⟩ :=
    if atEnd cfg i then r.val
    else let c := curr cfg i in
         if p c then r.val
         else takeUntilAux n (mkResultOk (next cfg i) ch st true)
  | ⟨Result.error _ _ _ _ _, h⟩ := unreachableError h

@[specialize] def takeUntil (p : Char → Bool) : parserM Unit :=
λ ps cfg r, takeUntilAux p cfg (inputSize cfg) r

def takeUntilNewLine : parserM Unit :=
takeUntil (= '\n')

def whitespace : parserM Unit :=
takeUntil (λ c, !c.isWhitespace)

-- setOption Trace.Compiler.boxed True
--- setOption pp.implicit True

def strAux (cfg : ParserConfig) (str : String) (error : String) : Nat → resultOk → pos → Result Unit
| 0     r j := mkError r error
| (n+1) r j :=
  if str.atEnd j then r.val
  else
    match r with
    | ⟨Result.ok _ i ch st e, _⟩ :=
      if atEnd cfg i then Result.error error i ch Syntax.missing true
      else if curr cfg i = str.get j then strAux n (mkResultOk (next cfg i) ch st true) (str.next j)
      else Result.error error i ch Syntax.missing true
    | ⟨Result.error _ _ _ _ _, h⟩ := unreachableError h

-- #exit

@[inline] def str (s : String) : parserM Unit :=
λ ps cfg r, strAux cfg s ("expected " ++ repr s) (inputSize cfg) r 0

@[specialize] def manyAux (p : parserM Unit) : Nat → Bool → parserM Unit
| 0     fst := pure ()
| (k+1) fst := λ ps cfg r,
  let i₀ := currPos r in
  let s₀ := currState r in
  match p ps cfg r with
  | Result.ok a i c s _    := manyAux k false ps cfg (mkResultOk i c s)
  | Result.error _ _ c _ _ := Result.ok () i₀ c s₀ fst

@[inline] def many (p : parserM Unit) : parserM Unit  :=
λ ps cfg r, manyAux p (inputSize cfg) true ps cfg r

@[inline] def many1 (p : parserM Unit) : parserM Unit  :=
p *> many p

def dummyParserCore : parserCore :=
λ cfg r, mkError r "dummy"

def testParser {α : Type} (x : parserM α) (input : String) : String :=
let r :=
  x { cmdParser := dummyParserCore, termParser := λ _, dummyParserCore }
    { filename := "test", input := input, FileMap := FileMap.fromString input, tokens := Lean.Parser.Trie.empty }
    (mkResultOk 0 {} {messages := MessageLog.empty}) in
match r with
| Result.ok _ i _ _ _      := "Ok at " ++ toString i
| Result.error msg i _ _ _ := "Error at " ++ toString i ++ ": " ++ msg

/-
mutual def recCmd, recTerm (parseCmd : Parser) (parseTerm : Nat → Parser) (parseLvl : Nat → parserCore)
with recCmd  : Nat → parserCore
| 0     cfg r := mkError r "Parser: no progress"
| (n+1) cfg r := parseCmd ⟨recCmd n, parseLvl, recTerm n⟩ cfg r
with recTerm : Nat → Nat → parserCore
| 0     rbp cfg r := mkError r "Parser: no progress"
| (n+1) rbp cfg r := parseTerm rbp ⟨recCmd n, parseLvl, recTerm n⟩ cfg r
-/

/-
def runParser (x : Parser) (parseCmd : Parser) (parseLvl : Nat → Parser) (parseTerm : Nat → Parser)
               (input : Iterator) (cfg : ParserConfig) : Result Syntax :=
let it := input in
let n  := it.remaining in
let r  := mkResultOk it {} {messages := MessageLog.Empty} in
let pl := recLvl (parseLvl) n in
let ps : recParsers := { cmdParser  := recCmd parseCmd parseTerm pl n,
                          lvlParser  := pl,
                          termParser := recTerm parseCmd parseTerm pl n } in
x ps cfg r
-/

structure parsingTables :=
(leadingTermParsers : TokenMap Parser)
(trailingTermParsers : TokenMap trailingParser)

abbreviation CommandParserM (α : Type) :=
parsingTables → parserM α

end flatParser
end Lean

section
open Lean.flatParser

def flatP : parserM Unit :=
many1 (str "++" <|> str "**" <|>  (str "--" *> takeUntil (= '\n') *> any *> pure ()))

end

section
open Lean.Parser
open Lean.Parser.MonadParsec

@[reducible] def Parser (α : Type) : Type :=  ReaderT Lean.flatParser.recParsers (ReaderT Lean.flatParser.ParserConfig (ParsecT Syntax (StateT ParserCache Id))) α

def testParsec (p : Parser Unit) (input : String) : String :=
let ps : Lean.flatParser.recParsers := { cmdParser := Lean.flatParser.dummyParserCore, termParser := λ _, Lean.flatParser.dummyParserCore } in
let cfg : Lean.flatParser.ParserConfig := { filename := "test", input := input, FileMap := Lean.flatParser.FileMap.fromString input, tokens := Lean.Parser.Trie.empty } in
let r := p ps cfg input.mkOldIterator {} in
match r with
| (Parsec.Result.ok _ it _, _)   := "OK at " ++ toString it.offset
| (Parsec.Result.error msg _, _) := "Error " ++ msg.toString

@[inline] def str' (s : String) : Parser Unit :=
str s *> pure ()

def parsecP : Parser Unit :=
many1' (str' "++" <|> str' "**" <|> (str "--" *> takeUntil (λ c, c = '\n') *> any *> pure ()))

def parsecP2 : Parser Unit :=
many1' ((parseStringLiteral *> whitespace *> pure ()) <|> (str "--" *> takeUntil (λ c, c = '\n') *> any *> pure ()))

end

namespace BasicParser
open Lean.Parser
open Lean.Parser.MonadParsec

def testBasicParser (p : BasicParserM Unit) (input : String) : String :=
let cfg : Lean.Parser.ParserConfig := {
filename := "test", input := input, fileMap := { lines := {} }, tokens := Lean.Parser.Trie.empty } in
let r := p cfg input.mkOldIterator {} in
match r with
| (Parsec.Result.ok _ it _, _)   := "OK at " ++ toString it.offset
| (Parsec.Result.error msg _, _) := "Error " ++ msg.toString

@[inline] def str' (s : String) : BasicParserM Unit :=
str s *> pure ()

def parserP : BasicParserM Unit :=
many1' (str' "++" <|> str' "**" <|> (str "--" *> takeUntil (λ c, c = '\n') *> any *> pure ()))

def parser2 : BasicParserM Unit :=
many1' ((parseStringLiteral *> Lean.Parser.MonadParsec.whitespace *> pure ()) <|> (str "--" *> takeUntil (λ c, c = '\n') *> any *> pure ()))

def parser3 : BasicParserM Unit :=
Lean.Parser.whitespace

end BasicParser

def mkBigString : Nat → String → String
| 0     s := s
| (n+1) s := mkBigString n (s ++ "-- new comment\n")

def mkBigString2 : Nat → String → String
| 0     s := s
| (n+1) s := mkBigString2 n (s ++ "\"hello\\nworld\"\n-- comment\n")

def mkBigString3 : Nat → String → String
| 0     s := s
| (n+1) s := mkBigString3 n (s ++ "/- /- comment 1 -/ -/ \n -- comment 2 \n \t \n ")

@[noinline] def testFlatP (s : String) : IO Unit :=
IO.println (Lean.flatParser.testParser flatP s)

@[noinline] def testParsecP (p : Parser Unit) (s : String) : IO Unit :=
IO.println (testParsec p s)

@[noinline] def testBasicParser (p : Lean.Parser.BasicParserM Unit) (s : String) : IO Unit :=
IO.println (BasicParser.testBasicParser p s)

@[noinline] def prof {α : Type} (msg : String) (p : IO α) : IO α :=
let msg₁ := "Time for '" ++ msg ++ "':" in
let msg₂ := "Memory usage for '" ++ msg ++ "':" in
allocprof msg₂ (timeit msg₁ p)

def main (xs : List String) : IO Unit :=
-- let s₁ := mkBigString xs.head.toNat "" in
-- let s₂ := s₁ ++ "bad" ++ mkBigString 20 "" in
-- let s₃ := mkBigString2 xs.head.toNat "" in
let s₄ := mkBigString3 xs.head.toNat "" in
-- prof "flat Parser 1" (testFlatP s₁) *>
-- prof "flat Parser 2" (testFlatP s₂) *>
-- prof "Parsec 1" (testParsecP parsecP s₁) *>
-- prof "Parsec 2" (testParsecP parsecP s₂) *>
-- prof "Parsec 3" (testParsecP parsecP2 s₃) *>
prof "Basic parser 1" (testBasicParser BasicParser.parser3 s₄)


/-
Copyright (c) 2020 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
import Lean.Util.SCC
import Lean.Elab.PreDefinition.Basic
import Lean.Elab.PreDefinition.Structural
import Lean.Elab.PreDefinition.WF.Main
import Lean.Elab.PreDefinition.MkInhabitant

namespace Lean.Elab
open Meta
open Term

structure TerminationHints where
  terminationBy? : Option Syntax := none
  decreasingBy? : Option Syntax := none
  deriving Inhabited

private def addAndCompilePartial (preDefs : Array PreDefinition) (useSorry := false) : TermElabM Unit := do
  for preDef in preDefs do
    trace[Elab.definition] "processing {preDef.declName}"
    let all := preDefs.toList.map (·.declName)
    forallTelescope preDef.type fun xs type => do
      let value ← if useSorry then
        mkLambdaFVars xs (← mkSorry type (synthetic := true))
      else
        liftM <| mkInhabitantFor preDef.declName xs type
      addNonRec { preDef with
        kind  := DefKind.«opaque»
        value
      } (all := all)
  addAndCompilePartialRec preDefs

private def isNonRecursive (preDef : PreDefinition) : Bool :=
  Option.isNone $ preDef.value.find? fun
    | Expr.const declName _ => preDef.declName == declName
    | _ => false

private def partitionPreDefs (preDefs : Array PreDefinition) : Array (Array PreDefinition) :=
  let getPreDef    := fun declName => (preDefs.find? fun preDef => preDef.declName == declName).get!
  let vertices     := preDefs.toList.map (·.declName)
  let successorsOf := fun declName => (getPreDef declName).value.foldConsts [] fun declName successors =>
    if preDefs.any fun preDef => preDef.declName == declName then
      declName :: successors
    else
      successors
  let sccs := SCC.scc vertices successorsOf
  sccs.toArray.map fun scc => scc.toArray.map getPreDef

private def collectMVarsAtPreDef (preDef : PreDefinition) : StateRefT CollectMVars.State MetaM Unit := do
  collectMVars preDef.value
  collectMVars preDef.type

private def getMVarsAtPreDef (preDef : PreDefinition) : MetaM (Array MVarId) := do
  let (_, s) ← (collectMVarsAtPreDef preDef).run {}
  pure s.result

private def ensureNoUnassignedMVarsAtPreDef (preDef : PreDefinition) : TermElabM PreDefinition := do
  let pendingMVarIds ← getMVarsAtPreDef preDef
  if (← logUnassignedUsingErrorInfos pendingMVarIds) then
    let preDef := { preDef with value := (← mkSorry preDef.type (synthetic := true)) }
    if (← getMVarsAtPreDef preDef).isEmpty then
      return preDef
    else
      throwAbortCommand
  else
    return preDef

/--
  Letrec declarations produce terms of the form `(fun .. => ..) d` where `d` is a (partial) application of an auxiliary declaration for a letrec declaration.
  This method beta-reduces them to make sure they can be eliminated by the well-founded recursion module. -/
private def betaReduceLetRecApps (preDefs : Array PreDefinition) : MetaM (Array PreDefinition) :=
  preDefs.mapM fun preDef => do
    let value ← transform preDef.value fun e => do
      if e.isApp && e.getAppFn.isLambda && e.getAppArgs.all fun arg => arg.getAppFn.isConst && preDefs.any fun preDef => preDef.declName == arg.getAppFn.constName! then
        return .visit e.headBeta
      else
        return .continue
    return { preDef with value }

private def addAsAxioms (preDefs : Array PreDefinition) : TermElabM Unit := do
  for preDef in preDefs do
    let decl := Declaration.axiomDecl {
      name        := preDef.declName,
      levelParams := preDef.levelParams,
      type        := preDef.type,
      isUnsafe    := preDef.modifiers.isUnsafe
    }
    addDecl decl
    withSaveInfoContext do  -- save new env
      addTermInfo' preDef.ref (← mkConstWithLevelParams preDef.declName) (isBinder := true)
    applyAttributesOf #[preDef] AttributeApplicationTime.afterTypeChecking
    applyAttributesOf #[preDef] AttributeApplicationTime.afterCompilation

def addPreDefinitions (preDefs : Array PreDefinition) (hints : TerminationHints) : TermElabM Unit := withLCtx {} {} do
  for preDef in preDefs do
    trace[Elab.definition.body] "{preDef.declName} : {preDef.type} :=\n{preDef.value}"
  let preDefs ← preDefs.mapM ensureNoUnassignedMVarsAtPreDef
  let preDefs ← betaReduceLetRecApps preDefs
  let cliques := partitionPreDefs preDefs
  let mut terminationBy ← liftMacroM <| WF.expandTerminationBy hints.terminationBy? (cliques.map fun ds => ds.map (·.declName))
  let mut decreasingBy  ← liftMacroM <| WF.expandTerminationHint hints.decreasingBy? (cliques.map fun ds => ds.map (·.declName))
  let mut hasErrors := false
  for preDefs in cliques do
    trace[Elab.definition.scc] "{preDefs.map (·.declName)}"
    if preDefs.size == 1 && isNonRecursive preDefs[0]! then
      let preDef := preDefs[0]!
      if preDef.modifiers.isNoncomputable then
        addNonRec preDef
      else
        addAndCompileNonRec preDef
    else if preDefs.any (·.modifiers.isUnsafe) then
      addAndCompileUnsafe preDefs
    else if preDefs.any (·.modifiers.isPartial) then
      for preDef in preDefs do
        if preDef.modifiers.isPartial && !(← whnfD preDef.type).isForall then
          withRef preDef.ref <| throwError "invalid use of 'partial', '{preDef.declName}' is not a function{indentExpr preDef.type}"
      addAndCompilePartial preDefs
    else
      try
        let mut wf? := none
        let mut decrTactic? := none
        if let some wf := terminationBy.find? (preDefs.map (·.declName)) then
          wf? := some wf
          terminationBy := terminationBy.markAsUsed (preDefs.map (·.declName))
        if let some { ref, value := decrTactic } := decreasingBy.find? (preDefs.map (·.declName)) then
          decrTactic? := some (← withRef ref `(by $(⟨decrTactic⟩)))
          decreasingBy := decreasingBy.markAsUsed (preDefs.map (·.declName))
        if wf?.isSome || decrTactic?.isSome then
          wfRecursion preDefs wf? decrTactic?
        else
          withRef (preDefs[0]!.ref) <| mapError
            (orelseMergeErrors
              (structuralRecursion preDefs)
              (wfRecursion preDefs none none))
            (fun msg =>
              let preDefMsgs := preDefs.toList.map (MessageData.ofExpr $ mkConst ·.declName)
              m!"fail to show termination for{indentD (MessageData.joinSep preDefMsgs Format.line)}\nwith errors\n{msg}")
      catch ex =>
        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 if preDefs.all fun preDef => preDef.kind == DefKind.theorem then
            addAsAxioms preDefs
        catch _ => s.restore
  unless hasErrors do
    liftMacroM <| terminationBy.ensureAllUsed
    liftMacroM <| decreasingBy.ensureAllUsed

builtin_initialize
  registerTraceClass `Elab.definition.body
  registerTraceClass `Elab.definition.scc

end Lean.Elab


/-
Copyright (c) 2020 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura, Sebastian Ullrich
-/
import Lean.Util.CollectLevelParams
import Lean.Elab.DeclUtil
import Lean.Elab.DefView
import Lean.Elab.Inductive
import Lean.Elab.Structure
import Lean.Elab.MutualDef
import Lean.Elab.DeclarationRange
namespace Lean.Elab.Command

open Meta

private def ensureValidNamespace (name : Name) : MacroM Unit := do
  match name with
  | Name.str p s _ =>
    if s == "_root_" then
      Macro.throwError s!"invalid namespace '{name}', '_root_' is a reserved namespace"
    ensureValidNamespace p
  | Name.num p .. => Macro.throwError s!"invalid namespace '{name}', it must not contain numeric parts"
  | Name.anonymous => pure ()

/- Auxiliary function for `expandDeclNamespace?` -/
private def expandDeclIdNamespace? (declId : Syntax) : MacroM (Option (Name × Syntax)) := do
  let (id, optUnivDeclStx) := expandDeclIdCore declId
  let scpView := extractMacroScopes id
  match scpView.name with
  | Name.str Name.anonymous s _ => return none
  | Name.str pre s _            =>
    ensureValidNamespace pre
    let nameNew := { scpView with name := Name.mkSimple s }.review
    -- preserve "original" info, if any, so that hover etc. on the namespaced
    -- name access the info tree node of the declaration name
    let id := mkIdent nameNew |>.setInfo declId.getHeadInfo
    if declId.isIdent then
      return some (pre, id)
    else
      return some (pre, declId.setArg 0 id)
  | _ => return none

/- given declarations such as `@[...] def Foo.Bla.f ...` return `some (Foo.Bla, @[...] def f ...)` -/
private def expandDeclNamespace? (stx : Syntax) : MacroM (Option (Name × Syntax)) := do
  if !stx.isOfKind `Lean.Parser.Command.declaration then
    return none
  else
    let decl := stx[1]
    let k := decl.getKind
    if k == ``Lean.Parser.Command.abbrev ||
       k == ``Lean.Parser.Command.def ||
       k == ``Lean.Parser.Command.theorem ||
       k == ``Lean.Parser.Command.constant ||
       k == ``Lean.Parser.Command.axiom ||
       k == ``Lean.Parser.Command.inductive ||
       k == ``Lean.Parser.Command.classInductive ||
       k == ``Lean.Parser.Command.structure then
      match (← expandDeclIdNamespace? decl[1]) with
      | some (ns, declId) => return some (ns, stx.setArg 1 (decl.setArg 1 declId))
      | none              => return none
    else if k == ``Lean.Parser.Command.instance then
      let optDeclId := decl[3]
      if optDeclId.isNone then return none
      else match (← expandDeclIdNamespace? optDeclId[0]) with
        | some (ns, declId) => return some (ns, stx.setArg 1 (decl.setArg 3 (optDeclId.setArg 0 declId)))
        | none              => return none
    else
      return none

def elabAxiom (modifiers : Modifiers) (stx : Syntax) : CommandElabM Unit := do
  -- leading_parser "axiom " >> declId >> declSig
  let declId             := stx[1]
  let (binders, typeStx) := expandDeclSig stx[2]
  let scopeLevelNames ← getLevelNames
  let ⟨name, declName, allUserLevelNames⟩ ← expandDeclId declId modifiers
  addDeclarationRanges declName stx
  runTermElabM declName fun vars => Term.withLevelNames allUserLevelNames $ Term.elabBinders binders.getArgs fun xs => do
    Term.applyAttributesAt declName modifiers.attrs AttributeApplicationTime.beforeElaboration
    let type ← Term.elabType typeStx
    Term.synthesizeSyntheticMVarsNoPostponing
    let type ← instantiateMVars type
    let type ← mkForallFVars xs type
    let type ← mkForallFVars vars type (usedOnly := true)
    let (type, _) ← Term.levelMVarToParam type
    let usedParams  := collectLevelParams {} type |>.params
    match sortDeclLevelParams scopeLevelNames allUserLevelNames usedParams with
    | Except.error msg      => throwErrorAt stx msg
    | Except.ok levelParams =>
      let decl := Declaration.axiomDecl {
        name        := declName,
        levelParams := levelParams,
        type        := type,
        isUnsafe    := modifiers.isUnsafe
      }
      Term.ensureNoUnassignedMVars decl
      addDecl decl
      withSaveInfoContext do  -- save new env
        Term.addTermInfo declId (← mkConstWithLevelParams declName) (isBinder := true)
      Term.applyAttributesAt declName modifiers.attrs AttributeApplicationTime.afterTypeChecking
      if isExtern (← getEnv) declName then
        compileDecl decl
      Term.applyAttributesAt declName modifiers.attrs AttributeApplicationTime.afterCompilation

/-
leading_parser "inductive " >> declId >> optDeclSig >> optional ":=" >> many ctor
leading_parser atomic (group ("class " >> "inductive ")) >> declId >> optDeclSig >> optional ":=" >> many ctor >> optDeriving
-/
private def inductiveSyntaxToView (modifiers : Modifiers) (decl : Syntax) : CommandElabM InductiveView := do
  checkValidInductiveModifier modifiers
  let (binders, type?) := expandOptDeclSig decl[2]
  let declId           := decl[1]
  let ⟨name, declName, levelNames⟩ ← expandDeclId declId modifiers
  addDeclarationRanges declName decl
  let ctors      ← decl[4].getArgs.mapM fun ctor => withRef ctor do
    -- def ctor := leading_parser " | " >> declModifiers >> ident >> optional inferMod >> optDeclSig
    let ctorModifiers ← elabModifiers ctor[1]
    if ctorModifiers.isPrivate && modifiers.isPrivate then
      throwError "invalid 'private' constructor in a 'private' inductive datatype"
    if ctorModifiers.isProtected && modifiers.isPrivate then
      throwError "invalid 'protected' constructor in a 'private' inductive datatype"
    checkValidCtorModifier ctorModifiers
    let ctorName := ctor.getIdAt 2
    let ctorName := declName ++ ctorName
    let ctorName ← withRef ctor[2] $ applyVisibility ctorModifiers.visibility ctorName
    let inferMod := !ctor[3].isNone
    let (binders, type?) := expandOptDeclSig ctor[4]
    addDocString' ctorName ctorModifiers.docString?
    addAuxDeclarationRanges ctorName ctor ctor[2]
    pure { ref := ctor, modifiers := ctorModifiers, declName := ctorName, inferMod := inferMod, binders := binders, type? := type? : CtorView }
  let classes ← getOptDerivingClasses decl[5]
  pure {
    ref             := decl
    modifiers       := modifiers
    shortDeclName   := name
    declName        := declName
    levelNames      := levelNames
    binders         := binders
    type?           := type?
    ctors           := ctors
    derivingClasses := classes
  }

private def classInductiveSyntaxToView (modifiers : Modifiers) (decl : Syntax) : CommandElabM InductiveView :=
  inductiveSyntaxToView modifiers decl

def elabInductive (modifiers : Modifiers) (stx : Syntax) : CommandElabM Unit := do
  let v ← inductiveSyntaxToView modifiers stx
  elabInductiveViews #[v]

def elabClassInductive (modifiers : Modifiers) (stx : Syntax) : CommandElabM Unit := do
  let modifiers := modifiers.addAttribute { name := `class }
  let v ← classInductiveSyntaxToView modifiers stx
  elabInductiveViews #[v]

def getTerminationHints (stx : Syntax) : TerminationHints :=
  let decl := stx[1]
  let k := decl.getKind
  if k == ``Parser.Command.def || k == ``Parser.Command.theorem || k == ``Parser.Command.instance then
    let args := decl.getArgs
    { terminationBy? := args[args.size - 2].getOptional?, decreasingBy? := args[args.size - 1].getOptional? }
  else
    {}

@[builtinCommandElab declaration]
def elabDeclaration : CommandElab := fun stx => do
  match (← liftMacroM <| expandDeclNamespace? stx) with
  | some (ns, newStx) => do
    let ns := mkIdentFrom stx ns
    let newStx ← `(namespace $ns:ident $newStx end $ns:ident)
    withMacroExpansion stx newStx $ elabCommand newStx
  | none => do
    let modifiers ← elabModifiers stx[0]
    let decl     := stx[1]
    let declKind := decl.getKind
    if declKind == ``Lean.Parser.Command.«axiom» then
      elabAxiom modifiers decl
    else if declKind == ``Lean.Parser.Command.«inductive» then
      elabInductive modifiers decl
    else if declKind == ``Lean.Parser.Command.classInductive then
      elabClassInductive modifiers decl
    else if declKind == ``Lean.Parser.Command.«structure» then
      elabStructure modifiers decl
    else if isDefLike decl then
      elabMutualDef #[stx] (getTerminationHints stx)
    else
      throwError "unexpected declaration"

/- Return true if all elements of the mutual-block are inductive declarations. -/
private def isMutualInductive (stx : Syntax) : Bool :=
  stx[1].getArgs.all fun elem =>
    let decl     := elem[1]
    let declKind := decl.getKind
    declKind == `Lean.Parser.Command.inductive

private def elabMutualInductive (elems : Array Syntax) : CommandElabM Unit := do
  let views ← elems.mapM fun stx => do
     let modifiers ← elabModifiers stx[0]
     inductiveSyntaxToView modifiers stx[1]
  elabInductiveViews views

/- Return true if all elements of the mutual-block are definitions/theorems/abbrevs. -/
private def isMutualDef (stx : Syntax) : Bool :=
  stx[1].getArgs.all fun elem =>
    let decl := elem[1]
    isDefLike decl

private def isMutualPreambleCommand (stx : Syntax) : Bool :=
  let k := stx.getKind
  k == ``Lean.Parser.Command.variable ||
  k == ``Lean.Parser.Command.universe ||
  k == ``Lean.Parser.Command.check ||
  k == ``Lean.Parser.Command.set_option ||
  k == ``Lean.Parser.Command.open

private partial def splitMutualPreamble (elems : Array Syntax) : Option (Array Syntax × Array Syntax) :=
  let rec loop (i : Nat) : Option (Array Syntax × Array Syntax) :=
    if h : i < elems.size then
      let elem := elems.get ⟨i, h⟩
      if isMutualPreambleCommand elem then
        loop (i+1)
      else if i == 0 then
        none -- `mutual` block does not contain any preamble commands
      else
        some (elems[0:i], elems[i:elems.size])
    else
      none -- a `mutual` block containing only preamble commands is not a valid `mutual` block
  loop 0

@[builtinMacro Lean.Parser.Command.mutual]
def expandMutualNamespace : Macro := fun stx => do
  let mut ns?      := none
  let mut elemsNew := #[]
  for elem in stx[1].getArgs do
    match ns?, (← expandDeclNamespace? elem) with
    | _, none                         => elemsNew := elemsNew.push elem
    | none, some (ns, elem)           => ns? := some ns; elemsNew := elemsNew.push elem
    | some nsCurr, some (nsNew, elem) =>
      if nsCurr == nsNew then
        elemsNew := elemsNew.push elem
      else
        Macro.throwErrorAt elem s!"conflicting namespaces in mutual declaration, using namespace '{nsNew}', but used '{nsCurr}' in previous declaration"
  match ns? with
  | some ns =>
    let ns := mkIdentFrom stx ns
    let stxNew := stx.setArg 1 (mkNullNode elemsNew)
    `(namespace $ns:ident $stxNew end $ns:ident)
  | none => Macro.throwUnsupported

@[builtinMacro Lean.Parser.Command.mutual]
def expandMutualElement : Macro := fun stx => do
  let mut elemsNew := #[]
  let mut modified := false
  for elem in stx[1].getArgs do
    match (← expandMacro? elem) with
    | some elemNew => elemsNew := elemsNew.push elemNew; modified := true
    | none         => elemsNew := elemsNew.push elem
  if modified then
    pure $ stx.setArg 1 (mkNullNode elemsNew)
  else
    Macro.throwUnsupported

@[builtinMacro Lean.Parser.Command.mutual]
def expandMutualPreamble : Macro := fun stx =>
  match splitMutualPreamble stx[1].getArgs with
  | none => Macro.throwUnsupported
  | some (preamble, rest) => do
    let secCmd    ← `(section)
    let newMutual := stx.setArg 1 (mkNullNode rest)
    let endCmd    ← `(end)
    pure $ mkNullNode (#[secCmd] ++ preamble ++ #[newMutual] ++ #[endCmd])

@[builtinCommandElab «mutual»]
def elabMutual : CommandElab := fun stx => do
  let hints := { terminationBy? := stx[3].getOptional?, decreasingBy? := stx[4].getOptional? }
  if isMutualInductive stx then
    if let some bad := hints.terminationBy? then
      throwErrorAt bad "invalid 'termination_by' in mutually inductive datatype declaration"
    if let some bad := hints.decreasingBy? then
      throwErrorAt bad "invalid 'decreasing_by' in mutually inductive datatype declaration"
    elabMutualInductive stx[1].getArgs
  else if isMutualDef stx then
    for arg in stx[1].getArgs do
      let argHints := getTerminationHints arg
      if let some bad := argHints.terminationBy? then
        throwErrorAt bad "invalid 'termination_by' in 'mutual' block, it must be used after the 'end' keyword"
      if let some bad := argHints.decreasingBy? then
        throwErrorAt bad "invalid 'decreasing_by' in 'mutual' block, it must be used after the 'end' keyword"
    elabMutualDef stx[1].getArgs hints
  else
    throwError "invalid mutual block"

/- leading_parser "attribute " >> "[" >> sepBy1 (eraseAttr <|> Term.attrInstance) ", " >> "]" >> many1 ident -/
@[builtinCommandElab «attribute»] def elabAttr : CommandElab := fun stx => do
  let mut attrInsts := #[]
  let mut toErase := #[]
  for attrKindStx in stx[2].getSepArgs do
    if attrKindStx.getKind == ``Lean.Parser.Command.eraseAttr then
      let attrName := attrKindStx[1].getId.eraseMacroScopes
      unless isAttribute (← getEnv) attrName do
        throwError "unknown attribute [{attrName}]"
      toErase := toErase.push attrName
    else
      attrInsts := attrInsts.push attrKindStx
  let attrs ← elabAttrs attrInsts
  let idents := stx[4].getArgs
  for ident in idents do withRef ident <| liftTermElabM none do
    let declName ← resolveGlobalConstNoOverloadWithInfo ident
    Term.applyAttributes declName attrs
    for attrName in toErase do
      Attribute.erase declName attrName

def expandInitCmd (builtin : Bool) : Macro := fun stx => do
  let optVisibility := stx[0]
  let optHeader     := stx[2]
  let doSeq         := stx[3]
  let attrId        := mkIdentFrom stx $ if builtin then `builtinInit else `init
  if optHeader.isNone then
    unless optVisibility.isNone do
      Macro.throwError "invalid initialization command, 'visibility' modifer is not allowed"
    `(@[$attrId:ident]def initFn : IO Unit := do $doSeq)
  else
    let id   := optHeader[0]
    let type := optHeader[1][1]
    if optVisibility.isNone then
      `(def initFn : IO $type := do $doSeq
        @[$attrId:ident initFn] constant $id : $type)
    else if optVisibility[0].getKind == ``Parser.Command.private then
      `(def initFn : IO $type := do $doSeq
        @[$attrId:ident initFn] private constant $id : $type)
    else if optVisibility[0].getKind == ``Parser.Command.protected then
      `(def initFn : IO $type := do $doSeq
        @[$attrId:ident initFn] protected constant $id : $type)
    else
      Macro.throwError "unexpected visibility annotation"

@[builtinMacro Lean.Parser.Command.«initialize»] def expandInitialize : Macro :=
  expandInitCmd (builtin := false)

@[builtinMacro Lean.Parser.Command.«builtin_initialize»] def expandBuiltinInitialize : Macro :=
  expandInitCmd (builtin := true)

end Lean.Elab.Command


/-
Copyright (c) 2019 Simon Hudon. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Simon Hudon
-/
import control.traversable.basic
import tactic.simpa

setup_tactic_parser

private meta def loc.to_string_aux : option name → string
| none := "⊢"
| (some x) := to_string x

/-- pretty print a `loc` -/
meta def loc.to_string : loc → string
| (loc.ns []) := ""
| (loc.ns [none]) := ""
| (loc.ns ls) := string.join $ list.intersperse " " (" at" :: ls.map loc.to_string_aux)
| loc.wildcard := " at *"

/-- shift `pos` `n` columns to the left -/
meta def pos.move_left (p : pos) (n : ℕ) : pos :=
{ line := p.line, column := p.column - n }

namespace tactic

open list

/-- parse structure instance of the shape `{ field1 := value1, .. , field2 := value2 }` -/
meta def struct_inst : lean.parser pexpr :=
do tk "{",
   ls ← sep_by (skip_info (tk ","))
     ( sum.inl <$> (tk ".." *> texpr) <|>
       sum.inr <$> (prod.mk <$> ident <* tk ":=" <*> texpr)),
   tk "}",
   let (srcs,fields) := partition_map id ls,
   let (names,values) := unzip fields,
   pure $ pexpr.mk_structure_instance
     { field_names := names,
       field_values := values,
       sources := srcs }

/-- pretty print structure instance -/
meta def struct.to_tactic_format (e : pexpr) : tactic format :=
do r ← e.get_structure_instance_info,
   fs ← mzip_with (λ n v,
     do v ← to_expr v >>= pp,
        pure $ format!"{n} := {v}" )
     r.field_names r.field_values,
   let ss := r.sources.map (λ s, format!" .. {s}"),
   let x : format := format.join $ list.intersperse ", " (fs ++ ss),
   pure format!" {{{x}}"

/-- Attribute containing a table that accumulates multiple `squeeze_simp` suggestions -/
@[user_attribute]
private meta def squeeze_loc_attr :
  user_attribute unit (option (list (pos × string × list simp_arg_type × string))) :=
{ name := `_squeeze_loc,
  parser := fail "this attribute should not be used",
  descr := "table to accumulate multiple `squeeze_simp` suggestions" }

/-- dummy declaration used as target of `squeeze_loc` attribute -/
def squeeze_loc_attr_carrier := ()

run_cmd squeeze_loc_attr.set ``squeeze_loc_attr_carrier none tt

/-- Format a list of arguments for use with `simp` and friends. This omits the
list entirely if it is empty. -/
meta def render_simp_arg_list : list simp_arg_type → tactic format
| [] := pure ""
| args := (++) " " <$> to_line_wrap_format <$> args.mmap pp

/-- Emit a suggestion to the user. If inside a `squeeze_scope` block,
the suggestions emitted through `mk_suggestion` will be aggregated so that
every tactic that makes a suggestion can consider multiple execution of the
same invocation.
If `at_pos` is true, make the suggestion at `p` instead of the current position. -/
meta def mk_suggestion (p : pos) (pre post : string) (args : list simp_arg_type)
  (at_pos := ff) : tactic unit :=
do xs ← squeeze_loc_attr.get_param ``squeeze_loc_attr_carrier,
   match xs with
   | none := do
     args ← render_simp_arg_list args,
     if at_pos then
       @scope_trace _ p.line p.column $
         λ _, _root_.trace sformat!"{pre}{args}{post}" (pure () : tactic unit)
     else
       trace sformat!"{pre}{args}{post}"
   | some xs := do
     squeeze_loc_attr.set ``squeeze_loc_attr_carrier ((p,pre,args,post) :: xs) ff
   end

local postfix `?`:9001 := optional

/-- translate a `pexpr` into a `simp` configuration -/
meta def parse_config : option pexpr → tactic (simp_config_ext × format)
| none := pure ({}, "")
| (some cfg) :=
  do e ← to_expr ``(%%cfg : simp_config_ext),
     fmt ← has_to_tactic_format.to_tactic_format cfg,
     prod.mk <$> eval_expr simp_config_ext e
             <*> struct.to_tactic_format cfg

/-- translate a `pexpr` into a `dsimp` configuration -/
meta def parse_dsimp_config : option pexpr → tactic (dsimp_config × format)
| none := pure ({}, "")
| (some cfg) :=
  do e ← to_expr ``(%%cfg : simp_config_ext),
     fmt ← has_to_tactic_format.to_tactic_format cfg,
     prod.mk <$> eval_expr dsimp_config e
             <*> struct.to_tactic_format cfg

/-- `same_result proof tac` runs tactic `tac` and checks if the proof
produced by `tac` is equivalent to `proof`. -/
meta def same_result (pr : proof_state) (tac : tactic unit) : tactic bool :=
do s ← get_proof_state_after tac,
   pure $ some pr = s

private meta def filter_simp_set_aux
  (tac : bool → list simp_arg_type → tactic unit)
  (args : list simp_arg_type) (pr : proof_state) :
  list simp_arg_type → list simp_arg_type →
  list simp_arg_type → tactic (list simp_arg_type × list simp_arg_type)
| [] ys ds := pure (ys.reverse, ds.reverse)
| (x :: xs) ys ds :=
  do b ← same_result pr (tac tt (args ++ xs ++ ys)),
     if b
       then filter_simp_set_aux xs ys (x:: ds)
       else filter_simp_set_aux xs (x :: ys) ds

declare_trace squeeze.deleted

/--
`filter_simp_set g call_simp user_args simp_args` returns `args'` such that, when calling
`call_simp tt /- only -/ args'` on the goal `g` (`g` is a meta var) we end up in the same
state as if we had called `call_simp ff (user_args ++ simp_args)` and removing any one
element of `args'` changes the resulting proof.
-/
meta def filter_simp_set
  (tac : bool → list simp_arg_type → tactic unit)
  (user_args simp_args : list simp_arg_type) : tactic (list simp_arg_type) :=
do some s ← get_proof_state_after (tac ff (user_args ++ simp_args)),
   (simp_args', _)  ← filter_simp_set_aux tac user_args s simp_args [] [],
   (user_args', ds) ← filter_simp_set_aux tac simp_args' s user_args [] [],
   when (is_trace_enabled_for `squeeze.deleted = tt ∧ ¬ ds.empty)
     trace!"deleting provided arguments {ds}",
   pure (user_args' ++ simp_args')

/-- make a `simp_arg_type` that references the name given as an argument -/
meta def name.to_simp_args (n : name) : tactic simp_arg_type :=
do e ← resolve_name' n, pure $ simp_arg_type.expr e

/-- tactic combinator to create a `simp`-like tactic that minimizes its
argument list.

 * `slow`: adds all rfl-lemmas from the environment to the initial list (this is a slower but more
           accurate strategy)
 * `no_dflt`: did the user use the `only` keyword?
 * `args`:    list of `simp` arguments
 * `tac`:     how to invoke the underlying `simp` tactic

-/
meta def squeeze_simp_core
  (slow no_dflt : bool) (args : list simp_arg_type)
  (tac : Π (no_dflt : bool) (args : list simp_arg_type), tactic unit)
  (mk_suggestion : list simp_arg_type → tactic unit) : tactic unit :=
do v ← target >>= mk_meta_var,
   args ← if slow then do
     simp_set ← attribute.get_instances `simp,
     simp_set ← simp_set.mfilter $ has_attribute' `_refl_lemma,
     simp_set ← simp_set.mmap $ resolve_name' >=> pure ∘ simp_arg_type.expr,
     pure $ args ++ simp_set
   else pure args,
   g ← retrieve $ do
   { g ← main_goal,
     tac no_dflt args,
     instantiate_mvars g },
   let vs := g.list_constant,
   vs ← vs.mfilter is_simp_lemma,
   vs ← vs.mmap strip_prefix,
   vs ← vs.to_list.mmap name.to_simp_args,
   with_local_goals' [v] (filter_simp_set tac args vs)
     >>= mk_suggestion,
   tac no_dflt args

namespace interactive

attribute [derive decidable_eq] simp_arg_type

/-- Turn a `simp_arg_type` into a string. -/
meta instance simp_arg_type.has_to_string : has_to_string simp_arg_type :=
⟨λ a, match a with
| simp_arg_type.all_hyps := "*"
| (simp_arg_type.except n) := "-" ++ to_string n
| (simp_arg_type.expr e) := to_string e
| (simp_arg_type.symm_expr e) := "←" ++ to_string e
end⟩

/-- combinator meant to aggregate the suggestions issued by multiple calls
of `squeeze_simp` (due, for instance, to `;`).

Can be used as:

```lean
example {α β} (xs ys : list α) (f : α → β) :
  (xs ++ ys.tail).map f = xs.map f ∧ (xs.tail.map f).length = xs.length :=
begin
  have : xs = ys, admit,
  squeeze_scope
  { split; squeeze_simp,
    -- `squeeze_simp` is run twice, the first one requires
    -- `list.map_append` and the second one
    -- `[list.length_map, list.length_tail]`
    -- prints only one message and combine the suggestions:
    -- > Try this: simp only [list.length_map, list.length_tail, list.map_append]
    squeeze_simp [this]
    -- `squeeze_simp` is run only once
    -- prints:
    -- > Try this: simp only [this] },
end
```

-/
meta def squeeze_scope (tac : itactic) : tactic unit :=
do none ← squeeze_loc_attr.get_param ``squeeze_loc_attr_carrier | pure (),
   squeeze_loc_attr.set ``squeeze_loc_attr_carrier (some []) ff,
   finally tac $ do
     some xs ← squeeze_loc_attr.get_param ``squeeze_loc_attr_carrier | fail "invalid state",
     let m := native.rb_lmap.of_list xs,
     squeeze_loc_attr.set ``squeeze_loc_attr_carrier none ff,
     m.to_list.reverse.mmap' $ λ ⟨p,suggs⟩, do
       { let ⟨pre,_,post⟩ := suggs.head,
         let suggs : list (list simp_arg_type) := suggs.map $ prod.fst ∘ prod.snd,
         mk_suggestion p pre post (suggs.foldl list.union []) tt, pure () }

/--
`squeeze_simp`, `squeeze_simpa` and `squeeze_dsimp` perform the same
task with the difference that `squeeze_simp` relates to `simp` while
`squeeze_simpa` relates to `simpa` and `squeeze_dsimp` relates to
`dsimp`. The following applies to `squeeze_simp`, `squeeze_simpa` and
`squeeze_dsimp`.

`squeeze_simp` behaves like `simp` (including all its arguments)
and prints a `simp only` invocation to skip the search through the
`simp` lemma list.

For instance, the following is easily solved with `simp`:

```lean
example : 0 + 1 = 1 + 0 := by simp
```

To guide the proof search and speed it up, we may replace `simp`
with `squeeze_simp`:

```lean
example : 0 + 1 = 1 + 0 := by squeeze_simp
-- prints:
-- Try this: simp only [add_zero, eq_self_iff_true, zero_add]
```

`squeeze_simp` suggests a replacement which we can use instead of
`squeeze_simp`.

```lean
example : 0 + 1 = 1 + 0 := by simp only [add_zero, eq_self_iff_true, zero_add]
```

`squeeze_simp only` prints nothing as it already skips the `simp` list.

This tactic is useful for speeding up the compilation of a complete file.
Steps:

   1. search and replace ` simp` with ` squeeze_simp` (the space helps avoid the
      replacement of `simp` in `@[simp]`) throughout the file.
   2. Starting at the beginning of the file, go to each printout in turn, copy
      the suggestion in place of `squeeze_simp`.
   3. after all the suggestions were applied, search and replace `squeeze_simp` with
      `simp` to remove the occurrences of `squeeze_simp` that did not produce a suggestion.

Known limitation(s):
  * in cases where `squeeze_simp` is used after a `;` (e.g. `cases x; squeeze_simp`),
    `squeeze_simp` will produce as many suggestions as the number of goals it is applied to.
    It is likely that none of the suggestion is a good replacement but they can all be
    combined by concatenating their list of lemmas. `squeeze_scope` can be used to
    combine the suggestions: `by squeeze_scope { cases x; squeeze_simp }`
  * sometimes, `simp` lemmas are also `_refl_lemma` and they can be used without appearing in the
    resulting proof. `squeeze_simp` won't know to try that lemma unless it is called as
    `squeeze_simp?`
-/
meta def squeeze_simp
  (key : parse cur_pos)
  (slow_and_accurate : parse (tk "?")?)
  (use_iota_eqn : parse (tk "!")?) (no_dflt : parse only_flag) (hs : parse simp_arg_list)
  (attr_names : parse with_ident_list) (locat : parse location)
  (cfg : parse struct_inst?) : tactic unit :=
do (cfg',c) ← parse_config cfg,
   squeeze_simp_core slow_and_accurate.is_some no_dflt hs
     (λ l_no_dft l_args, simp use_iota_eqn none l_no_dft l_args attr_names locat cfg')
     (λ args,
        let use_iota_eqn := if use_iota_eqn.is_some then "!" else "",
            attrs := if attr_names.empty then ""
                     else string.join (list.intersperse " " (" with" :: attr_names.map to_string)),
            loc := loc.to_string locat in
        mk_suggestion (key.move_left 1)
          sformat!"Try this: simp{use_iota_eqn} only"
          sformat!"{attrs}{loc}{c}" args)

/-- see `squeeze_simp` -/
meta def squeeze_simpa
  (key : parse cur_pos)
  (slow_and_accurate : parse (tk "?")?)
  (use_iota_eqn : parse (tk "!")?) (no_dflt : parse only_flag) (hs : parse simp_arg_list)
  (attr_names : parse with_ident_list) (tgt : parse (tk "using" *> texpr)?)
  (cfg : parse struct_inst?) : tactic unit :=
do (cfg',c) ← parse_config cfg,
   tgt' ← traverse (λ t, do t ← to_expr t >>= pp,
                            pure format!" using {t}") tgt,
   squeeze_simp_core slow_and_accurate.is_some no_dflt hs
     (λ l_no_dft l_args, simpa use_iota_eqn none l_no_dft l_args attr_names tgt cfg')
     (λ args,
        let use_iota_eqn := if use_iota_eqn.is_some then "!" else "",
            attrs := if attr_names.empty then ""
                     else string.join (list.intersperse " " (" with" :: attr_names.map to_string)),
            tgt' := tgt'.get_or_else "" in
        mk_suggestion (key.move_left 1)
          sformat!"Try this: simpa{use_iota_eqn} only"
          sformat!"{attrs}{tgt'}{c}" args)

/-- `squeeze_dsimp` behaves like `dsimp` (including all its arguments)
and prints a `dsimp only` invocation to skip the search through the
`simp` lemma list. See the doc string of `squeeze_simp` for examples.
 -/
meta def squeeze_dsimp
  (key : parse cur_pos)
  (slow_and_accurate : parse (tk "?")?)
  (use_iota_eqn : parse (tk "!")?)
  (no_dflt : parse only_flag) (hs : parse simp_arg_list)
  (attr_names : parse with_ident_list) (locat : parse location)
  (cfg : parse struct_inst?) : tactic unit :=
do (cfg',c) ← parse_dsimp_config cfg,
   squeeze_simp_core slow_and_accurate.is_some no_dflt hs
     (λ l_no_dft l_args, dsimp l_no_dft l_args attr_names locat cfg')
     (λ args,
        let use_iota_eqn := if use_iota_eqn.is_some then "!" else "",
            attrs := if attr_names.empty then ""
                     else string.join (list.intersperse " " (" with" :: attr_names.map to_string)),
            loc := loc.to_string locat in
        mk_suggestion (key.move_left 1)
          sformat!"Try this: dsimp{use_iota_eqn} only"
          sformat!"{attrs}{loc}{c}" args)

end interactive
end tactic

open tactic.interactive
add_tactic_doc
{ name       := "squeeze_simp / squeeze_simpa / squeeze_dsimp / squeeze_scope",
  category   := doc_category.tactic,
  decl_names :=
   [``squeeze_simp,
    ``squeeze_dsimp,
    ``squeeze_simpa,
    ``squeeze_scope],
  tags       := ["simplification", "Try this"],
  inherit_description_from := ``squeeze_simp }