import Yatima.Typechecker.TypecheckM
import Yatima.Typechecker.Printing
import Yatima.Common.ToLDON
import Lurk.Scalar

/-!
# Yatima typechecker: Eval

## Basic Structure

This is the first of the three main files that constitute the Yatima typechecker: `Eval`, `Equal`,
and `Infer`.

TODO: Add a high level overview of Eval in the context of Eval-Equal-Infer.

## Evaluate

In this module the evaluation (↔ reduction) of Yatima expressions is defined. Expressions that can
be reduced take a few forms, for example `.app fnc args`, constants, and suspdended evaluations.
Functions that can not be reduced further evaluate to unreduced Values or suspended thunks waiting
to evaluate further.
-/

namespace Yatima

open IR
open Lurk (F)

namespace Typechecker

/--
Looks for a constant by its hash `f : F` in a store and
returns it if found. Panics otherwise.

In the code generator, this function has to be overwritten with `(open f)`,
ignoring the second argument.
-/
def derefConst (f : F) (store : Store) : Const :=
  store.find! f

/-- TODO document. This function is overwritten btw -/
def mkInductiveProjF (block : F) (idx : Nat) (quick : Bool) : F :=
  let indF : Const := .inductiveProj ⟨block, idx⟩
  if quick then .ofNat $ (Hashable.hash indF).toNat
  else indF.toLDON.commit default |>.1

/-- TODO document. This function is overwritten btw -/
def mkConstructorProjF (block : F) (idx : Nat) (cidx : Nat) (quick : Bool) : F :=
  let ctorF : Const := .constructorProj ⟨block, idx, cidx⟩
  if quick then .ofNat $ (Hashable.hash ctorF).toNat
  else ctorF.toLDON.commit default |>.1

/-- TODO document. This function is overwritten btw -/
def mkRecursorProjF (block : F) (idx : Nat) (ridx : Nat) (quick : Bool) : F :=
  let recrF : Const := .recursorProj ⟨block, idx, ridx⟩
  if quick then .ofNat $ (Hashable.hash recrF).toNat
  else recrF.toLDON.commit default |>.1

/-- TODO document. This function is overwritten btw -/
def mkDefinitionProjF (block : F) (idx : Nat) (quick : Bool) : F :=
  let defnF : Const := .definitionProj ⟨block, idx⟩
  if quick then .ofNat $ (Hashable.hash defnF).toNat
  else defnF.toLDON.commit default |>.1

/--
Looks for a constant by its hash `f : F` in the `TypecheckState` cache of `TypedConst` and
returns it if it is found. If the constant is not found it throws an error.
Specifically, this function assumes that `checkConst name f` has previously been called
(which populates this cache).

Note: The `name : Name` is used only in the error messaging
-/
def derefTypedConst (f : F) : TypecheckM TypedConst := do
  match (← get).typedConsts.find? f with
  | some const => pure const
  | none => throw s!"TypedConst for {f} not found"

end Typechecker

namespace IR

open Typechecker (TypecheckM derefConst)

def getIndFromProj : InductiveProj → TypecheckM Inductive
  | ⟨indBlockF, idx⟩ => do
    let .mutIndBlock inds := derefConst indBlockF (← read).store
      | throw "Invalid Const kind. Expected mutIndBlock"
    let some ind := inds.get? idx
      | throw s!"Mutual inductive block doesn't contain index {idx}"
    pure ind

def getDefFromProj : DefinitionProj → TypecheckM Definition
  | ⟨defBlockF, idx⟩ => do
    let .mutDefBlock defs := derefConst defBlockF (← read).store
      | throw "Invalid Const kind. Expected mutDefBlock"
    let some defn := defs.get? idx
      | throw s!"Mutual definition block doesn't contain index {idx}"
    pure defn

def getCtorFromProj : ConstructorProj → TypecheckM Constructor
  | ⟨indBlockF, idx, cidx⟩ => do
    let ind ← getIndFromProj ⟨indBlockF, idx⟩
    let some ctor := ind.ctors.get? cidx
      | throw s!"Inductive doesn't contain constructor with index {cidx}"
    pure ctor

def getRecrFromProj : RecursorProj → TypecheckM Recursor
  | ⟨indBlockF, idx, ridx⟩ => do
    let ind ← getIndFromProj ⟨indBlockF, idx⟩
    let some recr := ind.recrs.get? ridx
      | throw s!"Inductive doesn't contain recursor with index {ridx}"
    pure recr

namespace Const

def levels : Const → TypecheckM Nat
  | .axiom      x
  | .theorem    x
  | .opaque     x
  | .definition x
  | .quotient   x => pure x.lvls
  | .inductiveProj   p => do pure (← getIndFromProj  p).lvls
  | .constructorProj p => do pure (← getCtorFromProj p).lvls
  | .recursorProj    p => do pure (← getRecrFromProj p).lvls
  | .definitionProj  p => do pure (← getDefFromProj  p).lvls
  | _ => throw "Can't retrieve universe levels of mutual blocks"

def type : Const → TypecheckM Expr
  | .axiom      x
  | .theorem    x
  | .opaque     x
  | .definition x
  | .quotient   x => pure x.type
  | .inductiveProj   p => do pure (← getIndFromProj  p).type
  | .constructorProj p => do pure (← getCtorFromProj p).type
  | .recursorProj    p => do pure (← getRecrFromProj p).type
  | .definitionProj  p => do pure (← getDefFromProj  p).type
  | _ => throw "Can't retrieve type of mutual blocks"

end Const

end IR

namespace Typechecker

def TypeInfo.update (univs : List Univ) : TypeInfo → TypeInfo
| .sort lvl => .sort $ lvl.instBulkReduce univs
| .unit  => .unit
| .proof => .proof
| .none  => .none

open PP

mutual
  /--
  Evaluates a `TypedExpr` into a `Value`.

  Evaluation here means applying functions to arguments, resuming evaluation of suspended thunks,
  evaluating a constant, instantiating a universe variable, evaluating the body of a let binding
  and evaluating a projection.
  -/
  partial def eval (t : TypedExpr) : TypecheckM Value := match t.expr with
    | .app fnc arg => do
      let ctx ← read
      let argThunk := suspend arg ctx (← get)
      let fnc ← evalTyped fnc
      apply fnc argThunk
    | .lam dom bod => do
      let ctx ← read
      let dom' := suspend dom ctx (← get)
      pure $ .lam dom' bod ctx.env
    | .var idx => do
      let some thunk := (← read).env.exprs.get? idx
        | throw s!"Index {idx} is out of range for expression environment"
      pure $ thunk.get
    | .const f const_univs => do
      let env := (← read).env
      let const_univs := const_univs.map (Univ.instBulkReduce env.univs)
      evalConst f const_univs
    | .letE _ val bod => do
      let thunk := suspend val (← read) (← get)
      withExtendedEnv thunk (eval bod)
    | .pi dom img => do
      let ctx ← read
      let dom' := suspend dom ctx (← get)
      pure $ .pi dom' img ctx.env
    | .sort univ => do
      let env := (← read).env
      pure $ .sort (Univ.instBulkReduce env.univs univ)
    | .lit lit =>
      pure $ .lit lit
    | .proj ind idx expr => do
      let val ← eval expr
      match val with
      | .app (.const f _) args _ =>
        match derefConst f (← read).store with
        | .constructorProj p =>
          let ctor ← getCtorFromProj p
          -- Since terms are well-typed, we can be sure that this constructor is of a structure-like inductive
          -- and, furthermore, that the index is in range of `args`
          let idx := ctor.params + idx
          let some arg := args.reverse.get? idx
            | throw s!"Invalid projection of index {idx} but constructor has only {args.length} arguments"
          pure $ arg.get
        | _ => pure $ .neu (.proj ind idx (.mk (expr.info.update (← read).env.univs) val))
      | .app .. => pure $ .neu (.proj ind idx (.mk (expr.info.update (← read).env.univs) val))
      | e => throw s!"Value {← ppValue e} is impossible to project"

  @[inline]
  partial def evalTyped (t : TypedExpr) : TypecheckM TypedValue := do
    let reducedInfo := t.info.update (← read).env.univs
    let value ← eval t
    pure ⟨reducedInfo, value⟩

  partial def evalConst' (f : F) (univs : List Univ) : TypecheckM Value := do
    match derefConst f (← read).store with
    | .theorem _
    | .definition _ =>
      match ← derefTypedConst f with
      | .theorem _ deref => withEnv ⟨[], univs⟩ $ eval deref
      | .definition _ deref part =>
        if part then pure $ mkConst f univs
        else withEnv ⟨[], univs⟩ $ eval deref
      | _ => throw "Invalid const kind for evaluation"
    | _ => pure $ mkConst f univs

  /-- Evaluates the `Yatima.Const` that's referenced by a constant index -/
  partial def evalConst (const : F) (univs : List Univ) : TypecheckM Value := do
    if ← primFWith .natZero (pure false) (pure $ · == const) then pure $ .lit (.natVal 0)
    else if (← fPrim const) matches .some (.op _) then pure $ mkConst const univs
    else evalConst' const univs

  /--
  Suspends the evaluation of a Yatima expression `expr : TypedExpr` in a particular `ctx : TypecheckCtx`

  Suspended evaluations can be resumed by evaluating `Thunk.get` on the resulting Thunk.
  -/
  partial def suspend (expr : TypedExpr) (ctx : TypecheckCtx) (stt : TypecheckState) : SusValue :=
    let thunk := { fn := fun _ =>
      match TypecheckM.run ctx stt (eval expr) with
      | .ok a =>
        a
      | .error e => .exception e }
    let reducedInfo := expr.info.update ctx.env.univs
    ⟨reducedInfo, thunk⟩

  /--
  Applies `value : Value` to the argument `arg : SusValue`.

  Applications are split into cases on whether `value` is a `Value.lam`, the application of a constant
  or the application of a free variable.

  * `Value.lam` : Descends into and evaluates the body of the lambda expression
  * `Value.app (.const ..)` : Applies the constant to the argument as expected using `applyConst`
  * `Value.app (.fvar ..)` : Returns an unevaluated `Value.app`
  -/
  partial def apply (val : TypedValue) (arg : SusValue) : TypecheckM Value :=
    match val.value with
    | .lam _ bod lamEnv =>
      withNewExtendedEnv lamEnv arg (eval bod)
    | .app (.const f kUnivs) args infos => applyConst f kUnivs arg args val.info infos
    -- Note that `val.info` is being added to the `infos` field of the `app` nodes because it is the info
    -- of the former partial application. That's because a stuck application like `h a1 .. an a(n+1)` must
    -- hold the info of the sub stuck application `h a1 .. an` for quoting to be done correctly
    | .app neu args infos => pure $ .app neu (arg :: args) (val.info :: infos)
    -- Since terms are well-typed we know that any other case is impossible
    | _ => throw "Invalid case for apply"

  /--
  Applies a named constant, referred by its constant index `f : F` to the list of arguments
  `arg :: args`.

  The application of the constant is split into cases on whether it is an inductive recursor,
  a quotient, or any other constant (which returns an unreduced application)
   -/
  partial def applyConst (f : F) (univs : List Univ) (arg : SusValue) (args : List SusValue)
      (info : TypeInfo) (infos : List TypeInfo) : TypecheckM Value := do
    if let some $ .op p ← fPrim f then
      if args.length < p.numArgs - 1 then
        pure $ .app (.const f univs) (arg :: args) (info :: infos)
      else
        let op := p.toPrimOp
        let argsArr := (Array.mk $ arg :: args).reverse
        match ← op.op $ argsArr with
        | .some v => pure v
        | .none =>
          if p.reducible then
            let typArgs := (info :: infos).zip (arg :: args)
            typArgs.foldrM (init := ← evalConst' f univs)
              fun (info, arg) acc => apply ⟨info, acc⟩ arg
          else pure $ .app (.const f univs) (arg :: args) (info :: infos)

    -- Assumes a partial application of f to args, which means in particular,
    -- that it is in normal form
    else match ← derefTypedConst f with
    | .recursor _ params motives minors indices isK indProj rules =>
      let majorIdx := params + motives + minors + indices
      if args.length != majorIdx then
        pure $ .app (.const f univs) (arg :: args) (info :: infos)
      else if isK then
        -- sanity check
        let nArgs := args.length
        let nDrop := params + motives + 1
        if nArgs < nDrop then
          throw s!"Too few arguments ({nArgs}). At least {nDrop} needed"
        let minorIdx := nArgs - nDrop
        let some minor := args.get? minorIdx | throw s!"Index {minorIdx} is out of range"
        pure minor.get
      else
        let params := args.take params
        match ← toCtorIfLitOrStruct indProj params univs arg with
        | .app (Neutral.const f _) args' _ => match ← derefTypedConst f with
          | .constructor _ idx _ =>
            match rules.get? idx with
            | some (fields, rhs) =>
              let exprs := (args'.take fields) ++ (args.drop indices)
              withEnv ⟨exprs, univs⟩ $ eval rhs.toImplicitLambda
            -- Since we assume expressions are previously type checked, we know that this constructor
            -- must have an associated recursion rule
            | none => throw s!"Constructor {f} has no associated recursion rule"
          | _ => pure $ .app (Neutral.const f univs) (arg :: args) (info :: infos)
        | _ => pure $ .app (Neutral.const f univs) (arg :: args) (info :: infos)
    | .quotient _ kind => match kind with
      | .lift => applyQuot arg args 6 1 (.app (.const f univs) (arg :: args) (info :: infos))
      | .ind  => applyQuot arg args 5 0 (.app (.const f univs) (arg :: args) (info :: infos))
      | _ => pure $ .app (.const f univs) (arg :: args) (info :: infos)
    | _ => pure $ .app (.const f univs) (arg :: args) (info :: infos)

  /--
  Applies a quotient to a value. It might reduce if enough arguments are applied to it
  -/
  partial def applyQuot (major? : SusValue) (args : List SusValue)
      (reduceSize argPos : Nat) (default : Value) : TypecheckM Value :=
    let argsLength := args.length + 1
    if argsLength == reduceSize then
      match major?.get with
      | .app (.const majorFn _) majorArgs _ => do
        match ← derefTypedConst majorFn with
        | .quotient _ .ctor =>
          -- Sanity check (`majorArgs` should have size 3 if the typechecking is correct)
          if majorArgs.length != 3 then throw "majorArgs should have size 3"
          let some majorArg := majorArgs.head? | throw "majorArgs can't be empty"
          let some head := args.get? argPos | throw s!"{argPos} is an invalid index for args"
          apply head.getTyped majorArg
        | _ => pure default
      | _ => pure default
    else if argsLength < reduceSize then
      pure default
    else
      throw s!"argsLength {argsLength} can't be greater than reduceSize {reduceSize}"

  partial def toCtorIfLitOrStruct (indProj : InductiveProj) (params : List SusValue) (univs : List Univ) : SusValue → TypecheckM Value
    | .mk info thunk =>
      match thunk.get with
      | .lit (.natVal v) => do
        let zeroIdx ← primF .natZero
        let succIdx ← primF (.op .natSucc)
        if v == 0 then pure $ mkConst zeroIdx []
        else
          let thunk : SusValue := ⟨info, Value.lit $ .natVal (v-1)⟩
          pure $ .app (.const succIdx []) [thunk] [.none]
      | .lit (.strVal _) => throw "TODO Reduction of string"
      | e => do
        -- do not eta expand structs in `Prop`
        if info == .proof then return e
        else match indProj with
        | ⟨f, i⟩ =>
          let ind ← getIndFromProj indProj
          -- must be a struct to eta expand
          if !ind.struct then
            pure e
          else
            let quick := (← read).quick
            let ctorF := mkConstructorProjF f i 0 quick
            match e with
            | .app (.const f _) _ _ => if ctorF == f then
              -- already eta expanded
              return e
            | _ => pure ()
            let ctor ← match ind.ctors with
              | [ctor] => pure ctor
              | _ =>
                let f := mkInductiveProjF f i (← read).quick
                throw s!"{(← read).constNames.getF f} should be a struct with only one constructor"
            let etaExpand (e : Value) : TypecheckM Value := do
              let mut projArgs : List SusValue := params
              for idx in [:ctor.fields] do
                -- FIXME get the correct TypeInfo for the projection
                projArgs := projArgs ++ [.mk .none $ .mk fun _ =>
                  .neu (.proj (mkInductiveProjF f i quick) idx $ .mk info e)]
              let len := projArgs.length
              if h : len > 0 then
                let lastIdx := len.pred
                let lastArg := projArgs.get ⟨lastIdx, Nat.pred_lt' h⟩
                let annotatedArgs := projArgs.take lastIdx ++ [lastArg]
                pure $ .app (.const ctorF univs) annotatedArgs $ annotatedArgs.map (fun _ => .none)
              else
                pure $ .neu (.const ctorF univs)
            etaExpand e
end

mutual
  /--
  Quoting transforms a value into a (typed) expression. It is the right-inverse of evaluation:
  evaluating a quoted value results in the value itself.
  -/
  partial def quote (lvl : Nat) (env : Env) : Value → TypecheckM Expr
    | .sort univ => pure $ .sort (univ.instBulkReduce env.univs)
    | .app neu args infos => do
      -- Sanity check: `args` and `infos` should have the same size
      if args.tail.length != infos.tail.length then throw "Partial application does not have enough info"
      let argsInfos := args.zip infos
      argsInfos.foldrM (init := ← quoteNeutral lvl env neu) fun (arg, info) acc => do
        pure $ .app ⟨info, acc⟩ $ ← quoteTyped lvl env arg.getTyped
    | .lam dom bod env' => do
      let dom ← quoteTyped lvl env dom.getTyped
      -- NOTE: although we add a value with `default` as `TypeInfo`, this is overwritten by the info of the expression's value
      let var := mkSusVar default lvl
      let bod ← quoteTypedExpr (lvl+1) bod (env'.extendWith var)
      pure $ .lam dom bod
    | .pi dom img env' => do
      let dom ← quoteTyped lvl env dom.getTyped
      let var := mkSusVar default lvl
      let img ← quoteTypedExpr (lvl+1) img (env'.extendWith var)
      pure $ .pi dom img
    | .lit lit => pure $ .lit lit
    | .exception e => throw e

  @[inline]
  partial def quoteTyped (lvl : Nat) (env : Env) (val : TypedValue) : TypecheckM TypedExpr := do
    pure ⟨val.info, ← quote lvl env val.value⟩

  partial def quoteExpr (lvl : Nat) (expr : Expr) (env : Env) : TypecheckM Expr :=
    match expr with
    | .var idx => do
      match env.exprs.get? idx with
      -- NOTE: if everything is correct, then `info` should coincide with `val.info`. We will choose `info` since
      -- this allows us to add values to the environment without knowing which `TypeInfo` it should take. See their
      -- previous note
     | some val => quote lvl env val.get
     | none => throw s!"Unbound variable _@{idx}"
    | .app fnc arg => do
      let fnc ← quoteTypedExpr lvl fnc env
      let arg ← quoteTypedExpr lvl arg env
      pure $ .app fnc arg
    | .lam dom bod => do
      let dom ← quoteTypedExpr lvl dom env
      let var := mkSusVar default lvl
      let bod ← quoteTypedExpr (lvl+1) bod (env.extendWith var)
      pure $ .lam dom bod
    | .letE typ val bod => do
      let typ ← quoteTypedExpr lvl typ env
      let val ← quoteTypedExpr lvl val env
      let var := mkSusVar default lvl
      let bod ← quoteTypedExpr (lvl+1) bod (env.extendWith var)
      pure $ .letE typ val bod
    | .pi dom img => do
      let dom ← quoteTypedExpr lvl dom env
      let var := mkSusVar default lvl
      let img ← quoteTypedExpr (lvl+1) img (env.extendWith var)
      pure $ .pi dom img
    | .proj ind idx expr => do
      let expr ← quoteTypedExpr lvl expr env
      pure $ .proj ind idx expr
    | .const idx univs => pure $ .const idx (univs.map (Univ.instBulkReduce env.univs))
    | .sort univ => pure $ .sort (univ.instBulkReduce env.univs)
    | .lit .. => pure expr

  @[inline]
  partial def quoteTypedExpr (lvl : Nat) (t : TypedExpr) (env : Env) : TypecheckM TypedExpr := do
    pure ⟨t.info, ← quoteExpr lvl t.expr env⟩

  partial def quoteNeutral (lvl : Nat) (env : Env) : Neutral → TypecheckM Expr
    | .fvar  idx => pure $ .var (lvl - idx - 1)
    | .const cidx univs => pure $ .const cidx (univs.map (Univ.instBulkReduce env.univs))
    | .proj  f ind val => do
      pure $ .proj f ind (← quoteTyped lvl env val)
end

end Yatima.Typechecker


import Lib.Meta.Opaque

namespace Bar
namespace Foo

opaque def Term := Unit

end Foo

namespace Foo

opaque namespace Term

open Lean

def fvar (x : Name) : Term :=
()

open Lean.Parser.Transport

def rec {motive : Term → Sort u} (f : ∀ n, motive (fvar n)) : ∀ t, motive t :=
sorry

theorem foo (n) : fvar n = fvar n := sorry

#check @Foo.Term.rec

end Term

#print Term.foo

end Foo


import data.list.perm
import data.multiset.basic

mk_iff_of_inductive_prop list.chain test.chain_iff

mk_iff_of_inductive_prop false    test.false_iff

mk_iff_of_inductive_prop true     test.true_iff

mk_iff_of_inductive_prop nonempty test.non_empty_iff

mk_iff_of_inductive_prop and      test.and_iff

mk_iff_of_inductive_prop or       test.or_iff

mk_iff_of_inductive_prop eq       test.eq_iff

mk_iff_of_inductive_prop heq      test.heq_iff

mk_iff_of_inductive_prop list.perm  test.perm_iff

mk_iff_of_inductive_prop list.pairwise  test.pairwise_iff

inductive test.is_true (p : Prop) : Prop
| triviality : p → test.is_true

mk_iff_of_inductive_prop test.is_true test.is_true_iff

@[mk_iff] structure foo (m n : ℕ) : Prop :=
(equal : m = n)
(sum_eq_two : m + n = 2)

example (m n : ℕ) : foo m n ↔ m = n ∧ m + n = 2 := foo_iff m n

@[mk_iff bar] structure foo2 (m n : ℕ) : Prop :=
(equal : m = n)
(sum_eq_two : m + n = 2)

example (m n : ℕ) : foo2 m n ↔ m = n ∧ m + n = 2 := bar m n


/-
Copyright (c) 2018 Johannes Hölzl. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Author: Johannes Hölzl
-/
import Mathlib.PrePort
import Mathlib.Lean3Lib.init.default
import Mathlib.data.finset.basic
import Mathlib.data.multiset.pi
import Mathlib.PostPort

universes u_1 u_2

namespace Mathlib

/-!
# The cartesian product of finsets
-/

namespace finset

/-! ### pi -/

/-- The empty dependent product function, defined on the empty set. The assumption `a ∈ ∅` is never
satisfied. -/
def pi.empty {α : Type u_1} (β : α → Type u_2) (a : α) (h : a ∈ ∅) : β a := multiset.pi.empty β a h

/-- Given a finset `s` of `α` and for all `a : α` a finset `t a` of `δ a`, then one can define the
finset `s.pi t` of all functions defined on elements of `s` taking values in `t a` for `a ∈ s`.
Note that the elements of `s.pi t` are only partially defined, on `s`. -/
def pi {α : Type u_1} {δ : α → Type u_2} [DecidableEq α] (s : finset α)
    (t : (a : α) → finset (δ a)) : finset ((a : α) → a ∈ s → δ a) :=
  mk (multiset.pi (val s) fun (a : α) => val (t a)) sorry

@[simp] theorem pi_val {α : Type u_1} {δ : α → Type u_2} [DecidableEq α] (s : finset α)
    (t : (a : α) → finset (δ a)) : val (pi s t) = multiset.pi (val s) fun (a : α) => val (t a) :=
  rfl

@[simp] theorem mem_pi {α : Type u_1} {δ : α → Type u_2} [DecidableEq α] {s : finset α}
    {t : (a : α) → finset (δ a)} {f : (a : α) → a ∈ s → δ a} :
    f ∈ pi s t ↔ ∀ (a : α) (h : a ∈ s), f a h ∈ t a :=
  multiset.mem_pi (val s) (fun (a : α) => (fun (a : α) => val (t a)) a) f

/-- Given a function `f` defined on a finset `s`, define a new function on the finset `s ∪ {a}`,
equal to `f` on `s` and sending `a` to a given value `b`. This function is denoted
`s.pi.cons a b f`. If `a` already belongs to `s`, the new function takes the value `b` at `a`
anyway. -/
def pi.cons {α : Type u_1} {δ : α → Type u_2} [DecidableEq α] (s : finset α) (a : α) (b : δ a)
    (f : (a : α) → a ∈ s → δ a) (a' : α) (h : a' ∈ insert a s) : δ a' :=
  multiset.pi.cons (val s) a b f a' sorry

@[simp] theorem pi.cons_same {α : Type u_1} {δ : α → Type u_2} [DecidableEq α] (s : finset α)
    (a : α) (b : δ a) (f : (a : α) → a ∈ s → δ a) (h : a ∈ insert a s) : pi.cons s a b f a h = b :=
  multiset.pi.cons_same (pi.cons._proof_1 s a a h)

theorem pi.cons_ne {α : Type u_1} {δ : α → Type u_2} [DecidableEq α] {s : finset α} {a : α} {a' : α}
    {b : δ a} {f : (a : α) → a ∈ s → δ a} {h : a' ∈ insert a s} (ha : a ≠ a') :
    pi.cons s a b f a' h = f a' (or.resolve_left (iff.mp mem_insert h) (ne.symm ha)) :=
  multiset.pi.cons_ne (pi.cons._proof_1 s a a' h) (ne.symm ha)

theorem pi_cons_injective {α : Type u_1} {δ : α → Type u_2} [DecidableEq α] {a : α} {b : δ a}
    {s : finset α} (hs : ¬a ∈ s) : function.injective (pi.cons s a b) :=
  sorry

@[simp] theorem pi_empty {α : Type u_1} {δ : α → Type u_2} [DecidableEq α]
    {t : (a : α) → finset (δ a)} : pi ∅ t = singleton (pi.empty δ) :=
  rfl

@[simp] theorem pi_insert {α : Type u_1} {δ : α → Type u_2} [DecidableEq α]
    [(a : α) → DecidableEq (δ a)] {s : finset α} {t : (a : α) → finset (δ a)} {a : α}
    (ha : ¬a ∈ s) :
    pi (insert a s) t = finset.bUnion (t a) fun (b : δ a) => image (pi.cons s a b) (pi s t) :=
  sorry

theorem pi_singletons {α : Type u_1} [DecidableEq α] {β : Type u_2} (s : finset α) (f : α → β) :
    (pi s fun (a : α) => singleton (f a)) = singleton fun (a : α) (_x : a ∈ s) => f a :=
  sorry

theorem pi_const_singleton {α : Type u_1} [DecidableEq α] {β : Type u_2} (s : finset α) (i : β) :
    (pi s fun (_x : α) => singleton i) = singleton fun (_x : α) (_x : _x ∈ s) => i :=
  pi_singletons s fun (_x : α) => i

theorem pi_subset {α : Type u_1} {δ : α → Type u_2} [DecidableEq α] {s : finset α}
    (t₁ : (a : α) → finset (δ a)) (t₂ : (a : α) → finset (δ a))
    (h : ∀ (a : α), a ∈ s → t₁ a ⊆ t₂ a) : pi s t₁ ⊆ pi s t₂ :=
  fun (g : (a : α) → a ∈ s → δ a) (hg : g ∈ pi s t₁) =>
    iff.mpr mem_pi fun (a : α) (ha : a ∈ s) => h a ha (iff.mp mem_pi hg a ha)

theorem pi_disjoint_of_disjoint {α : Type u_1} [DecidableEq α] {δ : α → Type u_2}
    [(a : α) → DecidableEq (δ a)] {s : finset α} [DecidableEq ((a : α) → a ∈ s → δ a)]
    (t₁ : (a : α) → finset (δ a)) (t₂ : (a : α) → finset (δ a)) {a : α} (ha : a ∈ s)
    (h : disjoint (t₁ a) (t₂ a)) : disjoint (pi s t₁) (pi s t₂) :=
  sorry

end Mathlib


/-
Copyright (c) 2017 Scott Morrison. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Scott Morrison
-/
import Mathlib.PrePort
import Mathlib.Lean3Lib.init.default
import Mathlib.tactic.auto_cases
import Mathlib.tactic.chain
import Mathlib.tactic.norm_cast
import Mathlib.PostPort

namespace Mathlib

namespace tactic

namespace tidy

/-- Tag interactive tactics (locally) with `[tidy]` to add them to the list of default tactics
called by `tidy`. -/
end tidy

namespace interactive

/-- Use a variety of conservative tactics to solve goals.

`tidy?` reports back the tactic script it found. As an example
```lean
example : ∀ x : unit, x = unit.star :=
begin
  tidy? -- Prints the trace message: "Try this: intros x, exact dec_trivial"
end
```

The default list of tactics is stored in `tactic.tidy.default_tidy_tactics`.
This list can be overridden using `tidy { tactics := ... }`.
(The list must be a `list` of `tactic string`, so that `tidy?`
can report a usable tactic script.)

Tactics can also be added to the list by tagging them (locally) with the
`[tidy]` attribute. -/
end interactive

/-- Invoking the hole command `tidy` ("Use `tidy` to complete the goal") runs the tactic of
the same name, replacing the hole with the tactic script `tidy` produces.
-/


import tactic
import data.list.alist
import control.monad.basic

def disableClangFormat := ff

class has_hmul (α β : Type*) (γ : out_param Type*) :=
  (mul : α → β → γ)
instance hmul_of_mul {α : Type*} [has_mul α] : has_hmul α α α := ⟨has_mul.mul⟩
infix ` ⋆ `:71 := has_hmul.mul

@[reducible] def Ident := string
@[reducible] def Label := string

@[derive [decidable_eq, fintype]]
inductive BinOp | add | sub | mul | lt | eq | lit_eq | and | or | min | max

/-- Expressions for a simple imperative language. -/
@[derive decidable_eq]
inductive E
| lit     : ℕ → E
| ident   : Ident → E
| not     : E → E
| bin_op  : BinOp → E → E → E
| access  : E → E → E
| ternary : E → E → E → E
| inline_code : string → E

-- todo simplify
| incr  : E → E
| attr  : E → E → E
| arrow_op : E → E → E
| call0 : E → E
| call1 : E → E → E
| call2 : E → E → E → E

/-- Statements for a simple imperative language, including sequencing. -/
inductive Prog
| skip
| accum (dst : E) (val : E)
| store (dst : E) (val : E)
| «if» (b : E) (cons : Prog) (alt : Prog)
| seq (a b : Prog)

| block (body : Prog)
| time  (n : string) (body : Prog)

| while (cond : E) (body : Prog)
| for (var bound : E) (body : Prog)

-- todo simplify
| declare (var : E) (initializer : E)
| auto (var : E) (initializer : E)

| inline_code (code : string)
| expr : E → Prog
| comment (c : string)

--def Prog.accum := λ l r, Prog.store l (l + r)

@[pattern]
def Prog.if1 (b : E) (cons : Prog) : Prog := Prog.if b cons Prog.skip

@[pattern] def E.false : E := E.lit 0
@[pattern] def E.true  : E := E.lit 1

namespace E

def neg : E → E
| (E.true) := E.false
| (E.false) := E.true
| e := e.not

infixr ` <;> `:1 := Prog.seq
infixr ` ;; `:1 := Prog.seq
infixr (name := seq) ` ; `:1 := Prog.seq
--instance : has_andthen Prog Prog Prog := ⟨Prog.seq⟩

instance : has_zero E := ⟨E.lit 0⟩
instance : has_one E  := ⟨E.lit 1⟩
instance : inhabited E := ⟨0⟩

instance : has_coe string E := ⟨E.ident⟩
end E

def BinOp.mk_type : BinOp → Type
| _ := E → E → E

-- a little smart
def BinOp.mk : Π (b : BinOp), BinOp.mk_type b
| BinOp.and := λ x y,
  match x, y with
  | E.true, y := y
  | x, E.true := x
  | E.false, y := E.false
  | x, E.false := E.false
  | x, y := E.bin_op BinOp.and x y
  end
| BinOp.or := λ x y,
  match x, y with
  | E.false, y := y
  | x, E.false := x
  | E.true, _ := E.true
  | _, E.true := E.true
  | x, y := E.bin_op BinOp.or x y
  end
| BinOp.add := λ x y,
  match x, y with
  | E.lit a, E.lit b := E.lit (a+b)
  | E.lit 0, x := x
  | x, E.lit 0 := x
  | _, _ := E.bin_op BinOp.add x y
  end
| BinOp.mul := λ x y,
  match x, y with
  | E.lit a, E.lit b := E.lit (a*b)
  | E.lit 0, x := E.lit 0
  | x, E.lit 0 := E.lit 0
  | E.lit 1, x := x
  | x, E.lit 1 := x
  | _, _ := E.bin_op BinOp.mul x y
  end
| b := E.bin_op b

instance : has_coe_to_fun BinOp BinOp.mk_type := ⟨BinOp.mk⟩

instance : has_add E := ⟨BinOp.add⟩
instance : has_sub E := ⟨BinOp.sub⟩
instance : has_mul E := ⟨BinOp.mul⟩

namespace E
def store   : E → E → Prog := Prog.store
def accum   : E → E → Prog := Prog.accum
def declare : E → E → Prog := Prog.declare
end E

section codegen

inductive ValueType | int | float
open ValueType
inductive TensorType
| atom (ty : ValueType)
| storage (ty : TensorType)
| sparse (ty : TensorType)

open TensorType

inductive CType
| double | int | storage (ty : CType) | sparse (ty : CType)

namespace TensorType
def toCType : TensorType → CType
| (atom ValueType.int) := CType.int
| (atom ValueType.float) := CType.double
| (storage t) := CType.storage (toCType t)
| (sparse t) := CType.sparse (toCType t)
end TensorType

@[reducible] def SymbolTable := alist (λ (s : string), TensorType)
structure Context :=
(true_conditions  : list E)
(false_conditions : list E)

def emptyContext := Context.mk [1] [0]

structure MState :=
(counter : ℕ)
(symbolTable : SymbolTable)
(buffer : buffer char)

instance : has_emptyc MState := ⟨⟨0, ∅, buffer.nil⟩⟩

-- note: inserting a writer_t for collecting output code kills performance
@[reducible] def M := state_t MState (reader Context)

def symbolType (var : string) : M TensorType :=
do
  s ← get,
  match s.symbolTable.lookup var with
  | some r := return r
  | none := return (atom ValueType.int) -- todo
  end

def runM {α} (m : M α) : α := ((m.run ∅).run emptyContext).fst
def BinOp.to_c : BinOp → string
| BinOp.add := "+"
| BinOp.sub := "-"
| BinOp.mul := "*"
| BinOp.lt := "int_lt"
| BinOp.eq := "int_eq"
| BinOp.lit_eq := "=="
| BinOp.and := "&&"
| BinOp.or := "||"
| BinOp.min := "min"
| BinOp.max := "max"

def wrap (s : string) : string := "(" ++ s ++ ")"

def E.to_c : E → string
| (E.lit i)                  := repr i
| (E.ident i)                := i
| (E.incr i)                 := wrap $ i.to_c ++ "++"
| (E.not i)                  := "!" ++ wrap i.to_c
| (E.bin_op BinOp.min e1 e2) := BinOp.min.to_c ++ (wrap $ e1.to_c ++ "," ++ e2.to_c)
| (E.bin_op BinOp.max e1 e2) := BinOp.max.to_c ++ (wrap $ e1.to_c ++ "," ++ e2.to_c)
| (E.bin_op BinOp.lt e1 e2)  := BinOp.lt.to_c ++ (wrap $ e1.to_c ++ "," ++ e2.to_c)
| (E.bin_op BinOp.eq e1 e2)  := BinOp.eq.to_c ++ (wrap $ e1.to_c ++ "," ++ e2.to_c)
| (E.bin_op op e1 e2)        := wrap $ e1.to_c ++ op.to_c ++ e2.to_c
| (E.ternary c t e)          := wrap $ wrap c.to_c ++ "?" ++ t.to_c ++ ":" ++ e.to_c
| (E.access e i)             := e.to_c ++ "[" ++ i.to_c ++ "]"
| (E.attr e i)               := e.to_c ++ "." ++ i.to_c
| (E.arrow_op e i)              := e.to_c ++ "->" ++ i.to_c
| (E.inline_code s)          := s
| (E.call0 f)                := f.to_c ++ wrap ""
| (E.call1 f a1)             := f.to_c ++ (wrap a1.to_c)
| (E.call2 f a1 a2)          := f.to_c ++ (wrap $ a1.to_c ++ "," ++ a2.to_c)

def emit (str : string) : M unit := modify $ λ s : MState,
{ s with buffer := s.buffer.append_string str }
def emitLine (s : string) : M unit := do emit $ s ++ ";"

namespace Prog

def to_c : Prog → M unit
| (expr e)          := emitLine $ e.to_c
| (accum dst val)   := emitLine $ dst.to_c ++ " += " ++ val.to_c
| (store dst val)   := emitLine $ dst.to_c ++ " = " ++ val.to_c
| (declare dst val) := emitLine $ "index " ++ dst.to_c ++ " = " ++ val.to_c
| (auto dst val)    := emitLine $ "auto " ++ dst.to_c ++ " = " ++ val.to_c
| (seq a b)         := a.to_c >> b.to_c
| (while c body)    := emit ("while" ++ wrap c.to_c ++ "{") >> body.to_c >> emit "}"
| (for i n body)    := emit ( "for" ++ wrap (i.to_c ++ "= 0;" ++ i.to_c ++ "<" ++ n.to_c ++ ";" ++ i.to_c++"++") ++ "{"
                          ) >> body.to_c >> emit "}"
| (inline_code s)   := emit s
| (skip)            := emit ""
| (comment s)       := emit $ "// " ++ s ++ "\n"
| (if1 E.false t)   := emit ""
| (if1 E.true t)    := t.to_c
| (if1 c t)         := do
    emit "if (" >> emit c.to_c >> emit ") {",
      t.to_c,
    emit "}"
| («if» c t e)      := do
    emit "if (" >> emit c.to_c >> emit ") {",
      t.to_c,
    emit "}", emit " else {",
      e.to_c,
    emit "}"
| (block p)         := emit "{" >> p.to_c >> emit "}"
| (time n p)          :=
  emit "{" >>
  ( emit ("cout << \"\\ntiming (" ++ n ++ "):\" << endl;") >>
    emit "out_val = 0.0; auto t1 = std::chrono::high_resolution_clock::now();" >>
    p.to_c >>
    emit "auto t2 = std::chrono::high_resolution_clock::now();" >>
    emit "cout << \"out: \" << out_val << endl;" >>
    emit "std::cout << \"took: \" << std::chrono::duration_cast<std::chrono::milliseconds>(t2-t1).count() << std::endl;"
  ) >> emit "}"
end Prog

def M.runInfo {α} (m : M α) : SymbolTable × buffer char :=
(λ s : MState, (s.symbolTable, s.buffer)) ((m.run ∅).run emptyContext).snd

def M.runBuffer {α} (m : M α) : string := m.runInfo.snd.to_string

def addHeaderFooter : string → string :=
λ s, "#include \"prefix.cpp\"\n" ++ s ++ "#include \"suffix.cpp\"\n"

def compile (progs : (list Prog)) : io unit :=
  let outName := "out_lean.cpp" in do
  handle ← io.mk_file_handle outName io.mode.write,
  let result : string := addHeaderFooter $ (progs.mmap (Prog.to_c ∘ Prog.block)).runBuffer,
  io.fs.write handle result.to_char_buffer,
  io.fs.close handle,
  if disableClangFormat then return () else io.cmd {cmd := "clang-format", args := ["-i", outName]} >> return ()

def comp : Prog → io unit := compile ∘ pure

-- #check monad.mapm

end codegen

section G

variables {α ι γ β : Type}

local infixl (name := and) ` && `:70 := BinOp.and
local infixl (name := or)  ` || `:65 := BinOp.or
local infix  (name := lt)  ` < `:71  := BinOp.lt
infix  (name := eq)  ` == `:71 := BinOp.eq
infix  (name := teq) ` === `:71 := BinOp.lit_eq
infix  (name := ne)  ` != `:71 := λ a b, (BinOp.eq a b).neg
@[pattern] def E.le : E → E → E := λ a b, BinOp.or (a < b) (a == b)
local infix  (name := le) ` ≤ `:71  := E.le

--notation e `⟦` k `⟧` := e.access k

structure G (ι α : Type) :=
  (index : ι)   (value : α)
  (ready : E)   (valid : E)
  (init : Prog) (next : Prog)

def G.empty [inhabited ι] [inhabited α] : G ι α :=
{ index := default, value := default, ready := E.false, valid := E.false, init := Prog.skip, next := Prog.skip }

structure View (ι α : Type) := (value : ι → α)

def constView (ι : Type) (v : α) : View ι α := ⟨λ _, v⟩
prefix (name := const) ` ⇑ ` := constView
-- instance : has_coe (α → β) (View α β) := ⟨View.mk⟩

instance {ι : Type*} : functor (G ι) :=
{ map := λ _ _ f g, { g with value := f g.value } }

instance {ι : Type*} : functor (View ι) :=
{ map := λ _ _ f v, { v with value := f ∘ v.value } }

instance View.has_hmul [has_hmul α β γ] : has_hmul (View ι α) (View ι β) (View ι γ) :=
⟨λ a b, ⟨λ i, a.value i ⋆ b.value i⟩⟩

namespace G
def iv {ι ι' α α'} (i : ι → ι') (v : α → α') : G ι α → G ι' α'
:= λ g, { g with value := v g.value, index := i g.index }

def mul [has_hmul α β γ] (a : G E α) (b : G E β) : G E γ :=
{ index := BinOp.max a.index b.index,
  value := a.value ⋆ b.value,
  ready := a.ready && b.ready && a.index == b.index,
  next  := Prog.if (a.index < b.index ||
                   (a.index == b.index && a.ready.neg))
                        a.next
                        b.next,
  valid := a.valid && b.valid,
  init  := a.init; b.init,
}
instance [has_hmul α β γ] : has_hmul (G E α) (G E β) (G E γ) := ⟨mul⟩

instance smul_G [has_smul E α] : has_smul E (G E α) :=
⟨λ s v, { v with value := s • v.value } ⟩
instance smul_unit [has_smul E α] : has_smul E (G unit α) :=
⟨λ s v, { v with value := s • v.value } ⟩
instance smul_base [has_hmul E α α] : has_smul E α := ⟨(⋆)⟩

example : has_smul E E := infer_instance

def add [has_smul E α] [has_mul α] [has_add α] (a b : G E α) : G E α :=
let current := BinOp.min a.index b.index in
{ index := current,
  value := (a.index == current) • a.value + (b.index == current) • b.value,
  ready := a.ready || b.ready,
  next  := Prog.if (a.index < b.index ||
                   (a.index == b.index && a.ready.neg && b.ready)) -- ($) <$>
             a.next
             $ (Prog.if (b.index < a.index ||
                        (a.index == b.index && b.ready.neg && a.ready))
                  b.next
                  (a.next; b.next)),
  valid := a.valid || b.valid,
  init  := a.init; b.init,
}

instance [has_mul α] [has_smul E α] [has_add α] : has_add (G E α) := ⟨add⟩

def mul_unit_const_r [has_hmul α β γ] (a : G unit α) (b : β) : G unit γ := (⋆ b) <$> a
def mul_unit_const_l [has_hmul α β γ] (a : α) (b : G unit β) : G unit γ := (λ v, a ⋆ v) <$> b
def mulViewR [has_hmul α β γ] (a : G ι α) (b : View ι β) : G ι γ :=
(⋆ b.value a.index) <$> a
def mulViewL [has_hmul α β γ] (a : View ι α) (b : G ι β) : G ι γ :=
(λ v, a.value b.index ⋆ v) <$> b

instance GV.has_hmul [has_hmul α β γ] : has_hmul (G ι α) (View ι β) (G ι γ) := ⟨G.mulViewR⟩
instance VG.has_hmul [has_hmul α β γ] : has_hmul (View ι α) (G ι β) (G ι γ) := ⟨G.mulViewL⟩
instance unit_const_r.has_hmul [has_hmul α β γ] : has_hmul (G unit α) β (G unit γ) := ⟨mul_unit_const_r⟩
instance unit_const_l.has_hmul [has_hmul α β γ] : has_hmul α (G unit β) (G unit γ) := ⟨mul_unit_const_l⟩
instance [has_mul α] : has_mul (G E α) := ⟨(⋆)⟩

end G

section simple_streams

def interval (i : E) (counter : E) (lower upper : E) : G E E :=
{ index := i.access counter, value := counter, ready := E.true, valid := counter < upper,
  init  := counter.declare lower, next  := counter.accum 1,
}

def range (counter bound : E) : G E E :=
{ index := counter, value := counter, ready := E.true, valid := counter < bound,
  init  := counter.declare 0, next  := counter.accum 1 }

def View.to_gen (counter bound : E) (view : View E α) : G E α := view.value <$> range counter bound

end simple_streams

section csr

/- implementation of composable sparse rval level -/
section rval

-- in TACO terminology, i = n_crd, v = n_pos. var indexes i.
structure csr := (i v var : E)

def csr.of (name : string) (n : ℕ) : csr :=
  let field (x : string) := E.ident $ name ++ n.repr ++ x in
  { i := field "_crd", v := field "_pos", var := field "_i" }

def csr'.of (name : string) (n : ℕ) : csr :=
  let field  (x : string) := E.ident $ name ++ n.repr ++ x in
  let field' (x : string) := E.ident $ name ++ (n+1).repr ++ x in
  { i := field "_crd", v := field' "_pos", var := field "_i" }

def csr.level : csr → E → G E E := λ csr loc,
interval csr.i csr.var (csr.v.access loc) (csr.v.access (loc+1))
def G.level   : csr → G E E → G E (G E E) := functor.map ∘ csr.level
def G.leaf    :   E → G E E → G E E       := functor.map ∘ E.access -- λ v, functor.map $ λ i, E.access v i

end rval

/- csr lval v3 -/
/- implementation of composable sparse lval level -/
section csr_lval
@[reducible] def loc := E
structure il :=
  (crd  : loc → E)
  (push : E → (loc → Prog) → Prog × loc)
structure vl  (α : Type) :=
  (pos  : loc → α)
  (init : loc → Prog)
structure lvl (α : Type) extends il, vl α.
instance : functor lvl := { map := λ _ _ f l, { l with pos := f ∘ l.pos } }

def sparse_index (indices : E) (bounds : E × E) : il :=
let upper := bounds.2, lower := bounds.1, current := indices.access (upper-1) in
let loc := upper-1 in
{ crd  := indices.access,
  push := λ i init,
    let prog := Prog.if1 (lower == upper || i != current)
                      ((upper.accum 1); init loc);
                     current.store i
    in (prog, loc) }

def dense_index (dim : E) (counter : E) (base : E) : il :=
{ crd  := id,
  push := λ i init,
    let l i  : loc := base * dim + i,
        prog : Prog := Prog.while (counter ≤ i) (init (l counter); counter.accum 1)
    in (prog, l i) }

def interval_vl (array : E) : vl (E × E) :=
let fn := array.access in
{ pos  := λ loc, (array.access loc, array.access (loc + 1)),
  init := λ loc, (fn (loc + 1)).store (fn loc) }

def dense_vl    (array : E) : vl E :=
{ pos := λ loc, array.access loc,
  init := λ loc, (array.access loc).store 0 }

def implicit_vl : vl E := { pos := id, init := λ _, Prog.skip }

-- def base (array : E) : lvl E := { i_shift := λ _ i, array.access i,  }

def rev_fmap_comp {f} [functor f] (x : α → f β) (y : β → f γ) := functor.map y ∘ x
infixr ` ⊚ `:90 := rev_fmap_comp
def rev_app : α → (α → β) → β := function.swap ($)
infixr ` & `:9 := rev_app

-- this combinator combines an il with a vl to form a lvl.
-- the extra parameter α is used to thread the primary argument to a level through ⊚.
--   see dcsr/csr_mat/dense below
def with_values : (α → il) → vl β → α → lvl β := λ i v e, lvl.mk (i e) v

def dense_mat (d₁ d₂ : E) := 0 &
  (with_values (dense_index d₁ "i1") implicit_vl) ⊚
  (with_values (dense_index d₂ "i2") $ dense_vl "values")

def cube_lvl := 0 &
  (with_values (dense_index 11 "i1") implicit_vl) ⊚
  (with_values (dense_index 7 "i2") implicit_vl) ⊚
  (with_values (dense_index 5 "i3") $ dense_vl "values")

def sparse_vec := (0, E.ident "size") &
  (with_values (sparse_index "A1_crd") (dense_vl "A_vals"))

def dcsr := (interval_vl "A1_pos").pos 0 &
  (with_values (sparse_index "A1_crd") (interval_vl "A2_pos")) ⊚
  (with_values (sparse_index "A2_crd") (dense_vl "A_vals"))

def csr_mat := 0 &
  (with_values (dense_index 2000 "i1") (interval_vl "B2_pos")) ⊚
  (with_values (sparse_index "B2_crd") (dense_vl "B_vals"))

end csr_lval

/- csr lval v2. TODO remove -/
structure lval (α : Type) := (push : E → Prog × α)
instance : functor lval := { map := λ _ _ f l, { l with push := prod.map id f ∘ l.push } }

def csr.base (csr : csr) : E×E := ((csr.v.access 0), (csr.v.access 1))
def csr.lval (csr : csr) (bounds : E×E) : lval (E×E)  :=
let upper := bounds.2, lower := bounds.1 in
{ push := λ i,
  let prog := Prog.if1 (lower == upper || (i != csr.i.access (upper-1) &&
                                            (csr.v.access upper != csr.v.access (upper-1))))
                 (upper.accum 1; (csr.v.access upper).store (csr.v.access (upper-1)));
               (csr.i.access (upper-1)).store i,
      value := ((csr.v.access (upper-1)), (csr.v.access upper)) in (prog, value)
}

def csr.vec (csr : csr) (bounds : E×E) : lval E :=
let upper := bounds.2, lower := bounds.1 in
{ push := λ i, (Prog.if1 (lower == upper || (i != csr.i.access (upper-1) &&
                                            (0 != csr.v.access (upper-1))))
                 (upper.accum 1; (csr.v.access (upper-1)).store 0);
               (csr.i.access (upper-1)).store i, csr.v.access (upper-1))
}

/- csr lval v1. TODO remove -/
structure pre  := (pre : Prog)
structure new  := (new : Prog) -- (acc : E → α)
structure post := (post : Prog)

-- push values at the leaf level
def push_value (var val_array outer_var : E) : new × post × E :=
({new := (val_array.access outer_var).store 0}, ⟨Prog.skip⟩, (val_array.access var))

-- if pack is true, we allow for an rval that produces duplicate coordinates and de-duplicate them as they are aggregated
def push_level_pack' (pack : bool) (csr : csr) (k : E → new × post × α) : (E → new × post × (E → pre × α))
:= λ outer_var,
( { new := (csr.v.access outer_var).store (csr.var+1) },
  { post := (csr.v.access (outer_var+1)).store (csr.var+1); (k csr.var).2.1.post },
   λ i, (⟨(if pack then Prog.if1 (csr.var < csr.v.access outer_var || i != csr.i.access csr.var) else id) $
             csr.var.accum 1; (csr.i.access csr.var).store i; (k csr.var).1.new⟩,
         (k csr.var).2.2))

def push_level_pack : csr → (E → new × post × α) → (E → new × post × (E → pre × α)) :=
push_level_pack' tt
def push_level      : csr → (E → new × post × α) → (E → new × post × (E → pre × α)) :=
push_level_pack' ff

def vec_lval' (n : string) :=
prod.snd $ (push_level_pack (csr.of n 1) $ push_value (csr.of n 1).var (E.ident $ n ++ "_vals")) 0
def mat_lval' (n : string) :=
prod.snd $ (push_level_pack (csr.of n 1) $ push_level_pack (csr.of n 2) $ push_value (csr.of n 2).var (E.ident $ n ++ "_vals")) 0
def mval (n : string) :=
prod.snd $ (push_level (csr.of n 1) $ push_level (csr.of n 2) $ push_value (csr.of n 2).var (E.ident $ n ++ "_vals")) 0
def cub_lval' (n : string) :=
prod.snd $ (push_level_pack (csr.of n 1) $ push_level_pack (csr.of n 2) $ push_level_pack (csr.of n 3) $ push_value (csr.of n 3).var (E.ident $ n ++ "_vals")) 0
def cub_lval'' (n : string) :=
prod.snd $ (push_level (csr.of n 1) $ push_level (csr.of n 2) $ push_level (csr.of n 3) $ push_value (csr.of n 3).var (E.ident $ n ++ "_vals")) 0
#check mat_lval'

end csr

def indexed_mat_lval (var : E) (i j v : E) := ((var.access i).access j).accum v

def fmap1 {α β f} [functor f] : (α → β) → f α → f β := functor.map
def fmap2 {α β f} [functor f] : (α → β) → f (f α) → f (f β) := functor.map ∘ functor.map

class Compile (l r : Type) := (eval : l → r → Prog)
class Scalar (α : Type) := (fold : α → E → Prog) (value : α → E)
instance : Scalar E := ⟨λ l r, l.accum r, id⟩

-- offset into array:
instance : Scalar (E × E) := ⟨λ l r, (l.2.access l.1).accum r, λ l, l.2.access l.1⟩

--instance discard.eval : Compile unit E := ⟨λ _ _, Prog.skip⟩
instance base.eval [Scalar α] : Compile α E :=
{ eval := λ l v, Scalar.fold l v }
-- todo remove
instance base.eval' : Compile (E → Prog) E :=
{ eval := λ acc v, acc v }

instance unit.eval [Compile α β] : Compile α (G unit β) :=
{ eval := λ acc v,
    v.init; Prog.while v.valid
      (Prog.if1 v.ready (Compile.eval acc v.value) ; v.next) }

instance lvl.eval [Compile α β] : Compile (lvl α) (G E β) :=
{ eval := λ storage v,
    let (push_i, loc) := storage.push v.index storage.init in
    v.init ;
    Prog.while v.valid
      (Prog.if1 v.ready
        (push_i;
         Compile.eval (storage.pos loc) v.value);
      v.next) }

instance lval.eval [Compile α β] : Compile (lval α) (G E β) :=
{ eval := λ acc v,
    let loop_body := acc.push v.index in
    v.init ;
    Prog.while v.valid
      (Prog.if1 v.ready
        (loop_body.1;
         Compile.eval loop_body.2 v.value);
      v.next) }

-- instance unit.bool [Scalar α] [Compile α β] : Compile α (G unit β) :=
-- { eval := λ acc v,
--     v.init; Prog.while (v.valid && (Scalar.value acc).neg)
--       (Prog.if1 v.ready (Compile.eval acc v.value) ; v.next) }

instance level.eval  [Compile α β] : Compile (E → α) (G E β) :=
{ eval := λ acc v,
    v.init; Prog.while v.valid
      (Prog.if1 v.ready (Compile.eval (acc v.index) v.value);
      v.next) }

instance level_pre.eval [Compile α β] : Compile (E → pre × α) (G E β) :=
{ eval := λ acc v,
    let loop_body := acc v.index in
    v.init ;
    Prog.while v.valid
      (Prog.if1 v.ready
        (loop_body.1.pre;
         Compile.eval loop_body.2 v.value);
      v.next) }

instance level_outer_post.eval [Compile α β] : Compile (post × α) β :=
{ eval := λ lhs v, let (x, acc) := lhs in Compile.eval acc v; x.post }

-- janky map-based generator:
-- i   = get<0>(x->first)
-- j   = get<1>(x->first)
-- val = x.second
def coo_rval (n : ℕ) (matrix var : E) : (G E E) :=
let call (a b : E) := E.call0 (a.attr b) in
{ index := E.call1 (E.ident $ "get<" ++ n.repr ++ ">") $ var.arrow_op "first",
  value := var.arrow_op "second",
  ready := E.true,
  valid := E.neg $ var === call matrix "end",
  init  := Prog.auto var $ call matrix "begin",
  next  := Prog.expr $ var.incr,
}
def coo_rval_inner (n : ℕ) (matrix var : E) : (G E E) :=
{ coo_rval n matrix var with
  next := Prog.inline_code "break;",
  valid := E.true,
  init := Prog.skip }

def coo_vector_rval (matrix var : E) : (G E E) := coo_rval 0 matrix var
def coo_matrix_rval (matrix var : E) : G E (G E E) :=
fmap1 (λ _, coo_rval_inner 1 matrix var) $ coo_vector_rval matrix var
def coo_cube_rval   (matrix var : E) : G E (G E (G E E)) :=
fmap2 (λ _, coo_rval_inner 2 matrix var) $ coo_matrix_rval matrix var

namespace G
def contract (g : G ι α) : G unit α := { g with index := () }
def sum1 : G E α → G unit α := contract
def sum2 : (G E (G E α)) → (G unit (G unit α)) :=
(functor.map $ sum1) ∘ contract
def sum2' : (G E (G E E)) → (G unit (G unit E)) :=
(functor.map $ sum1) ∘ contract
def sum3'  : G E (G E (G E α)) → G unit (G unit (G unit α)) :=
(functor.map $ sum2) ∘ contract
def sum3 : G E (G E (G E E)) → G unit (G unit (G unit E)) :=
(functor.map $ sum2) ∘ contract
def sum_inner : G E (G E (G E E)) → G E (G E (G unit E)) := functor.map $ functor.map G.contract

end G

def v  : G E E       := G.leaf "V_vals" $ ((csr.of "V" 1).level 0)
def A  : G E (G E E) := (csr.of "A" 1).level 0 & G.level (csr.of "A" 2) ⊚ G.leaf "A_vals"
def A_csr  : G E (G E E) :=
  let dense : G E E := range "_i" 2000 in
  dense & G.level (csr.of "A" 2) ⊚ G.leaf "A_vals"
def B_csr  : G E (G E E) :=
  let dense : G E E := range "_i" 2000 in
  dense & G.level (csr.of "B" 2) ⊚ G.leaf "B_vals"
def B  : G E (G E E) := (csr.of "B" 1).level 0 & G.level (csr.of "B" 2) ⊚ G.leaf "B_vals"
def C  : G E (G E (G E E)) := (csr.of "C" 1).level 0 & G.level (csr.of "C" 2) ⊚ G.level (csr.of "C" 3) ⊚ G.leaf "C_vals"
def D  : G E (G E (G E E)) := (csr.of "D" 1).level 0 & G.level (csr.of "D" 2) ⊚ G.level (csr.of "D" 3) ⊚ G.leaf "D_vals"
def M_  : G E (G E E) := G.leaf "M_vals" <$> ((csr.of "M" 1).level 0).level (csr.of "M" 2)

def exec {l r} [Compile l r] := @Compile.eval l r

def single (x : α) := (⇑ E x).to_gen "i" 1
#eval (single (3 : E)).ready.to_c
def eg00  := exec (E.ident "out") (E.lit 2)
def eg00' := exec (E.ident "out") (v.contract)
def eg01 := exec (Prog.accum "out") (G.contract <$> A.contract)
def eg02 := exec (Prog.accum "out") (G.sum2 $ A⋆B)
def eg03 := exec (Prog.accum "out") (G.sum2 $ (⇑ E v) ⋆ A)
def eg04 := exec (Prog.accum "out") (G.sum3 $ ((⇑ E) <$> A) ⋆ ⇑ E B ) -- (i,k)*(j,k)

def inner_prod := (fmap1 (⇑ E) A) ⋆ ⇑ E B -- (i,k)*(j,k)
-- (i,j)*(j,k)
def comb_rows : G E (G E (G unit E)) := G.sum_inner $ (fmap2 (⇑ E) A) ⋆ ⇑ E B

def me := Prog.time "me"
def ta := Prog.time "taco"

def c0 : csr := { i := "", v := "A1_pos", var := "" }
def c1 : csr := csr'.of "A" 1
def c2 : csr := { i := "A2_crd", v := "A_vals", var := "" }
def A_lval := c0.base & ( c1.lval ⊚ c2.vec )

def out : E := "out_val"

def eg05 := exec (Prog.accum "out") $ G.sum3 $ (fmap2 (⇑ E) A) ⋆ (fmap1 (⇑ E) B) -- (i,j)*(i,k)
def matsum := comb_rows.sum2
def eg06 := exec (Prog.accum "out") matsum
def eg07 := exec (indexed_mat_lval "out") (A⋆B)
def ref_vector (n : string) := (coo_vector_rval (E.ident $ n ++ ".data") "entry")
def ref_matrix (n : string) := (coo_matrix_rval (E.ident $ n ++ ".data") "entry")
def ref_cube   (n : string) := (coo_cube_rval   (E.ident $ n ++ ".data") "entry")
def ref_x := ref_matrix "x"
def ref_y := ref_matrix "y"
def ref_M := ref_matrix "M"
def eg18 := exec (λ (_ _ : E), Prog.accum "out") A
def load_AB' := [exec (mat_lval' "A") ref_x, exec (mat_lval' "B") ref_x]
def load_AB : list Prog := [
  Prog.time "gen A" $ exec (mat_lval' "A") (ref_matrix "x"),
  Prog.time "gen B" $ exec (mat_lval' "B") (ref_matrix "y") ]
def eg19 := load_AB' ++ [Prog.time "me" $ exec (mat_lval' "out") $ A⋆B]
def taco_ijk := Prog.inline_code "taco_ijk_sum();"
def eg21 := exec (vec_lval' "out") v
def eg22 := me $ exec (mat_lval' "out") A
def eg23 := load_AB' ++ [Prog.time "me" $ exec (mval "out") $ inner_prod.sum_inner]
def eg24 := exec (λ i, (Prog.accum $ out.access i)) ((λ (i : E), 2*i) <$> (range "i" 10))
def eg25 := exec (Prog.accum out) ((λ (i : E), 2*i) <$> (range "i" 10)).contract
def eg20 := load_AB ++ [me $ eg06, Prog.time "taco" $ taco_ijk]
def eg26 := load_AB ++ [me $ exec out $ G.sum2 (comb_rows), Prog.time "taco" $ taco_ijk]
def eg27 := load_AB ++ [me $ exec (mval "out") $ comb_rows]

def load := [
  exec A_lval (ref_matrix "x"),
  exec (mat_lval' "B") (ref_matrix "y"),
  exec (cub_lval' "C") (ref_cube "c"),
  exec (cub_lval' "D") (ref_cube "c"),
  exec (vec_lval' "V") (ref_vector "v")]

def eg28 := load_AB ++ [
  Prog.time "me" $ exec (mval "out") $ inner_prod.sum_inner,
  Prog.time "taco" $ Prog.inline_code "taco_ikjk();" ]

def eg29 := load ++ [exec (mval "out") (A + B)]
def eg30 := exec out (C + D).sum3

--#eval comp $ exec A_lval A
--#eval compile $ [exec A_lval (ref_matrix "x"), me $ exec out $ G.sum2 $ A ]
--#eval compile $ [exec sparse_vec (ref_vector "v")] -- , me $ exec out $ G.sum2 $ A ]
--#eval compile $ [exec dcsr (ref_matrix "x"), me $ exec out $ G.sum2 $ A ]

-- compare dcsr anc csr:
def eg31 := compile $
[ exec dcsr (ref_matrix "x"),
  exec csr_mat (ref_matrix "x"),
  me $ exec out $ B_csr.sum2,
  me $ exec out $ (ref_matrix "x").sum2 ]
--#eval eg31

end G


import classes.unrestricted.basics.toolbox
import utilities.list_utils

section functions_lift_sink

variables {T N₀ N : Type}

def lift_symbol_ (lift_N : N₀ → N) : symbol T N₀ → symbol T N
| (symbol.terminal t)    := symbol.terminal t
| (symbol.nonterminal n) := symbol.nonterminal (lift_N n)

def sink_symbol_ (sink_N : N → option N₀) : symbol T N → option (symbol T N₀)
| (symbol.terminal t)    := some (symbol.terminal t)
| (symbol.nonterminal n) := option.map symbol.nonterminal (sink_N n)

def lift_string_ (lift_N : N₀ → N) : list (symbol T N₀) → list (symbol T N) :=
list.map (lift_symbol_ lift_N)

def sink_string_ (sink_N : N → option N₀) : list (symbol T N) → list (symbol T N₀) :=
list.filter_map (sink_symbol_ sink_N)

def lift_rule_ (lift_N : N₀ → N) : grule T N₀ → grule T N :=
λ r : grule T N₀, grule.mk
  (lift_string_ lift_N r.input_L)
  (lift_N r.input_N)
  (lift_string_ lift_N r.input_R)
  (lift_string_ lift_N r.output_string)

end functions_lift_sink

section lifting_conditions

structure lifted_grammar_ (T : Type) :=
(g₀ g : grammar T)
(lift_nt : g₀.nt → g.nt)
(sink_nt : g.nt → option g₀.nt)
(lift_inj : function.injective lift_nt)
(sink_inj : ∀ x y, sink_nt x = sink_nt y →
  x = y  ∨  sink_nt x = none
)
(lift_nt_sink : ∀ n₀ : g₀.nt, sink_nt (lift_nt n₀) = some n₀)
(corresponding_rules : ∀ r : grule T g₀.nt,
  r ∈ g₀.rules →
    lift_rule_ lift_nt r ∈ g.rules
)
(preimage_of_rules : ∀ r : grule T g.nt,
  (r ∈ g.rules ∧ ∃ n₀ : g₀.nt, lift_nt n₀ = r.input_N) →
    (∃ r₀ ∈ g₀.rules, lift_rule_ lift_nt r₀ = r)
)

private lemma lifted_grammar_inverse {T : Type} (lg : lifted_grammar_ T) :
  ∀ x : lg.g.nt,
    (∃ val, lg.sink_nt x = some val) →
      option.map lg.lift_nt (lg.sink_nt x) = x :=
begin
  intros x h,
  cases h with valu ass,
  rw ass,
  rw option.map_some',
  apply congr_arg,
  symmetry,
  by_contradiction,
  have inje := lg.sink_inj x (lg.lift_nt valu),
  rw lg.lift_nt_sink at inje,
  cases inje ass with case_valu case_none,
  {
    exact h case_valu,
  },
  rw ass at case_none,
  exact option.no_confusion case_none,
end

end lifting_conditions

section translating_derivations

variables {T : Type}

private lemma lift_tran_ {lg : lifted_grammar_ T} {w₁ w₂ : list (symbol T lg.g₀.nt)}
    (hyp : grammar_transforms lg.g₀ w₁ w₂) :
  grammar_transforms lg.g (lift_string_ lg.lift_nt w₁) (lift_string_ lg.lift_nt w₂) :=
begin
  rcases hyp with ⟨r, rin, u, v, bef, aft⟩,
  use lift_rule_ lg.lift_nt r,
  split,
  {
    exact lg.corresponding_rules r rin,
  },
  use lift_string_ lg.lift_nt u,
  use lift_string_ lg.lift_nt v,
  split,
  {
    have lift_bef := congr_arg (lift_string_ lg.lift_nt) bef,
    unfold lift_string_ at *,
    rw list.map_append_append at lift_bef,
    rw list.map_append_append at lift_bef,
    exact lift_bef,
  },
  {
    have lift_aft := congr_arg (lift_string_ lg.lift_nt) aft,
    unfold lift_string_ at *,
    rw list.map_append_append at lift_aft,
    exact lift_aft,
  },
end

lemma lift_deri_ (lg : lifted_grammar_ T) {w₁ w₂ : list (symbol T lg.g₀.nt)}
    (hyp : grammar_derives lg.g₀ w₁ w₂) :
  grammar_derives lg.g (lift_string_ lg.lift_nt w₁) (lift_string_ lg.lift_nt w₂) :=
begin
  induction hyp with u v trash orig ih,
  {
    apply grammar_deri_self,
  },
  apply grammar_deri_of_deri_tran,
  {
    exact ih,
  },
  exact lift_tran_ orig,
end

def good_letter_ {lg : lifted_grammar_ T} : symbol T lg.g.nt → Prop
| (symbol.terminal t)    := true
| (symbol.nonterminal n) := (∃ n₀ : lg.g₀.nt, lg.sink_nt n = n₀)

def good_string_ {lg : lifted_grammar_ T} (s : list (symbol T lg.g.nt)) :=
∀ a ∈ s, good_letter_ a

private lemma sink_tran_ {lg : lifted_grammar_ T} {w₁ w₂ : list (symbol T lg.g.nt)}
    (hyp : grammar_transforms lg.g w₁ w₂)
    (ok_input : good_string_ w₁) :
  grammar_transforms lg.g₀ (sink_string_ lg.sink_nt w₁) (sink_string_ lg.sink_nt w₂)
  ∧ good_string_ w₂ :=
begin
  rcases hyp with ⟨r, rin, u, v, bef, aft⟩,

  rcases lg.preimage_of_rules r (by {
    split,
    {
      exact rin,
    },
    rw bef at ok_input,
    have good_matched_nonterminal : good_letter_ (symbol.nonterminal r.input_N),
    {
      apply ok_input (symbol.nonterminal r.input_N),
      apply list.mem_append_left,
      apply list.mem_append_left,
      apply list.mem_append_right,
      rw list.mem_singleton,
    },
    change ∃ n₀ : lg.g₀.nt, lg.sink_nt r.input_N = some n₀ at good_matched_nonterminal,
    cases good_matched_nonterminal with n₀ hn₀,
    use n₀,
    have almost := congr_arg (option.map lg.lift_nt) hn₀,
    rw lifted_grammar_inverse lg r.input_N ⟨n₀, hn₀⟩ at almost,
    rw option.map_some' at almost,
    apply option.some_injective,
    exact almost.symm,
  }) with ⟨r₀, pre_in, preimage⟩,

  split, swap,
  {
    rw bef at ok_input,
    rw aft,
    unfold good_string_ at ok_input ⊢,
    rw ←preimage,
    clear_except ok_input,
    rw list.forall_mem_append_append at ok_input ⊢,
    rw list.forall_mem_append_append at ok_input,
    split,
    {
      exact ok_input.1.1,
    },
    split, swap,
    {
      exact ok_input.2.2,
    },
    intros a a_in_ros,
    cases a,
    {
      clear_except,
      unfold good_letter_,
    },
    unfold lift_rule_ at a_in_ros,
    dsimp only at a_in_ros,
    unfold lift_string_ at a_in_ros,
    rw list.mem_map at a_in_ros,
    rcases a_in_ros with ⟨s, trash, a_from_s⟩,
    rw ←a_from_s,
    cases s,
    {
      exfalso,
      clear_except a_from_s,
      unfold lift_symbol_ at a_from_s,
      exact symbol.no_confusion a_from_s,
    },
    unfold lift_symbol_,
    unfold good_letter_,
    use s,
    exact lg.lift_nt_sink s,
  },
  use r₀,
  split,
  {
    exact pre_in,
  },
  use sink_string_ lg.sink_nt u,
  use sink_string_ lg.sink_nt v,
  have correct_inverse : sink_symbol_ lg.sink_nt ∘ lift_symbol_ lg.lift_nt = option.some,
  {
    ext1,
    cases x,
    {
      refl,
    },
    rw function.comp_app,
    unfold lift_symbol_,
    unfold sink_symbol_,
    rw lg.lift_nt_sink,
    apply option.map_some',
  },
  split,
  {
    have sink_bef := congr_arg (sink_string_ lg.sink_nt) bef,
    unfold sink_string_ at *,
    rw list.filter_map_append_append at sink_bef,
    rw list.filter_map_append_append at sink_bef,
    convert sink_bef;
    rw ←preimage;
    unfold lift_rule_;
    dsimp only;
    clear_except correct_inverse,
    {
      unfold lift_string_,
      rw list.filter_map_map,
      rw correct_inverse,
      rw list.filter_map_some,
    },
    {
      change
        [symbol.nonterminal r₀.input_N] =
        list.filter_map (sink_symbol_ lg.sink_nt)
          (list.map (lift_symbol_ lg.lift_nt) [symbol.nonterminal r₀.input_N]),
      rw list.filter_map_map,
      rw correct_inverse,
      rw list.filter_map_some,
    },
    {
      unfold lift_string_,
      rw list.filter_map_map,
      rw correct_inverse,
      rw list.filter_map_some,
    },
  },
  {
    have sink_aft := congr_arg (sink_string_ lg.sink_nt) aft,
    unfold sink_string_ at *,
    rw list.filter_map_append_append at sink_aft,
    convert sink_aft,
    rw ←preimage,
    clear_except correct_inverse,
    unfold lift_rule_,
    dsimp only,
    unfold lift_string_,
    rw list.filter_map_map,
    rw correct_inverse,
    rw list.filter_map_some,
  },
end

private lemma sink_deri_aux {lg : lifted_grammar_ T} {w₁ w₂ : list (symbol T lg.g.nt)}
    (hyp : grammar_derives lg.g w₁ w₂)
    (ok_input : good_string_ w₁) :
  grammar_derives lg.g₀ (sink_string_ lg.sink_nt w₁) (sink_string_ lg.sink_nt w₂)
  ∧ good_string_ w₂ :=
begin
  induction hyp with u v trash orig ih,
  {
    split,
    {
      apply grammar_deri_self,
    },
    {
      exact ok_input,
    },
  },
  have both := sink_tran_ orig ih.2,

  split, swap,
  {
    exact both.2,
  },
  apply grammar_deri_of_deri_tran,
  {
    exact ih.1,
  },
  {
    exact both.1,
  },
end

lemma sink_deri_ (lg : lifted_grammar_ T) {w₁ w₂ : list (symbol T lg.g.nt)}
    (hyp : grammar_derives lg.g w₁ w₂)
    (ok_input : good_string_ w₁) :
  grammar_derives lg.g₀ (sink_string_ lg.sink_nt w₁) (sink_string_ lg.sink_nt w₂) :=
begin
  exact (sink_deri_aux hyp ok_input).1
end

end translating_derivations


import Std

open Lean (Name FVarId getConstInfo MessageData MonadEnv)
open Lean.Meta
open Std.Tactic.Lint

namespace NamingConvention
  open Lean

  /-- If a name is part of the whitelist, then it won't be flagged as wrong even if it conflicts with
      the convention.

      Mostly used for automatically generated definitions that don't fit Mathlibs naming convention

-/
  def whitelist : List String := [
    "noConfusion",
    "noConfusionType",
    "injEq",
    "casesOn",
    "brecOn",
    "recOn"
  ]

  /-- Checks that `name` is in snake_case -/
  def snakeCaseTest (name : String) (reason : String) : MetaM (Option MessageData) := do
    let chars := name.toList
    let parts := (chars.splitOn '_').filter (fun x => !x.isEmpty)

    if parts.length == 1 then
      -- There are no `_`, so there should not be any upper case characters
      if chars.any Char.isUpper then
        return m!"`{name}` should be in snake_case ({reason})"
      else
        return none
    else
      let errors := List.reduceOption <| parts.map fun p =>
        if (p.get! 0).isUpper then
          some p
        else
          none

      if errors.isEmpty then
        return none
      else
        let errors := String.join <| errors.map fun error =>
          "\n * " ++ String.mk error ++ " should be in lowerCamelCase"
        return m!
          "When something named with UpperCamelCase is part of something named with snake_case,"
           ++ m!"it is referenced in lowerCamelCase ({reason}){errors}\n"

  /-- Checks whether the passed list contains a non-trailing underscore (`_`)
      An underscore is considered trailing if it is followed by only characters in [_,] until the end of the string
   -/
  def isSnakeCase : List Char → Bool
    | [] => false
    | '_' :: as => as.any (fun
                            | '_' | ',' => false
                            | _ => true
        )
    | _   :: as => isSnakeCase as

  /-- Checks that `name` is in lowerCamelCase, i.e., the first character is lowercase, and there
      are no `_` characters
   -/
  def lowerCamelCaseTest (name : String) (reason : String) : MetaM (Option MessageData) := do
    let chars := name.toList

    if (chars.get! 0).isUpper || isSnakeCase chars then
      return m!"`{name}` should be in lowerCamelCase ({reason})"
    else
      return none

  /-- Checks that `name` is in UpperCamelCase, i.e., the first character is Uppercase, and there
      are no `_` characters
    -/
  def upperCamelCaseTest (name : String) (reason : String) : MetaM (Option MessageData) := do
    let chars := name.toList

    if (chars.get! 0).isLower || isSnakeCase chars then
      return m!"`{name}` should be in UpperCamelCase ({reason})"
    else
      return none

  /--
    Checks that a `def` or `theorem` conforms to the naming convention

    https://github.com/leanprover-community/mathlib4/wiki/Porting-wiki#naming-convention
  -/
  def definitionTest (name : Name) (value : Expr) : MetaM (Option MessageData) := do
    if name.isInternal then
      return none

    -- Extract only the last element of the name
    let name ← match name with
      | .str _ name => pure name
      | _ => return none

    if whitelist.contains name then
      return none

    -- If a name starts with `inst` we assume it is a typeclass instance and don't flag it when misnamed
    -- TODO: proper detection when a definition is actually an instance, rather than relying on this heuristic
    if name.toList.take 4 == "inst".toList then
      return none

    let type ← try
      inferType value
    catch _ =>
      return none
    let type ← whnf type

    -- Functions are named the same way as their return values
    -- (e.g. a function of type A → B → C is named as though it is a term of type C
    forallTelescopeReducing type fun _ type => do
      if type.isSort then
        upperCamelCaseTest name "Props and Types (or Sort) (inductive types, structures, classes) are in UpperCamelCase"
      else
        let typeOfType ← inferType type

        if typeOfType.isProp then
          snakeCaseTest name "Terms of Props (e.g. proofs, theorem names) are in snake_case"
        else if typeOfType.isType then
          lowerCamelCaseTest name "Terms of Types (most definitions) are in lowerCamelCase"
        else
          pure none

end NamingConvention

open NamingConvention

/--
  Check whether the definition of the given name is consistent with the mathlib4 naming convention
-/
def namingConventionTest (name : Name) : MetaM (Option MessageData) := do
  let decl ← match (← MonadEnv.getEnv).find? name with
    | some decl => pure decl
    | none => return none

  match decl with
    | .defnInfo val => definitionTest name val.value
    | .thmInfo val  => definitionTest name val.value
    | _ => pure none

@[std_linter]
def namingConvention : Linter where
  test := namingConventionTest
  noErrorsFound := "No naming convention violations found"
  errorsFound := "Some definitions seem to be inconsistent with the naming convention"


import tidy.lib.list
import tidy.lib.env
import tidy.lib.tactic
import tidy.lib.pretty_print

open tactic tactic.interactive
open interactive

namespace tidy.rewrite_search.discovery

-- TODO make sure this didn't break anything
meta def is_acceptable_rewrite (t : expr) : bool :=
  is_eq_or_iff_after_binders t

meta def is_acceptable_lemma (r : expr) : tactic bool :=
  is_acceptable_rewrite <$> infer_type r

meta def is_acceptable_hyp (r : expr) : tactic bool :=
  do t ← infer_type r, return $ is_acceptable_rewrite t ∧ ¬t.has_meta_var

meta def assert_acceptable_lemma (r : expr) : tactic unit := do
  ret ← is_acceptable_lemma r,
  if ret then return ()
  else do
    pp ← pretty_print r,
    fail format!"\"{pp}\" is not a valid rewrite lemma!"

meta def load_attr_list : list name → tactic (list name)
| [] := return []
| (a :: rest) := do
  names ← attribute.get_instances a,
  l ← load_attr_list rest,
  return $ names ++ l

meta def load_names (l : list name) : tactic (list expr) :=
  l.mmap mk_const

meta def rewrite_list_from_rw_rules (rws : list rw_rule) : tactic (list (expr × bool)) :=
  rws.mmap (λ r, do e ← to_expr' r.rule, pure (e, r.symm))

meta def rewrite_list_from_lemmas (l : list expr) : list (expr × bool) :=
  l.map (λ e, (e, ff)) ++ l.map (λ e, (e, tt))

meta def rewrite_list_from_lemma (e : expr) : list (expr × bool) :=
  rewrite_list_from_lemmas [e]

meta def rewrite_list_from_hyps : tactic (list (expr × bool)) := do
  hyps ← local_context,
  rewrite_list_from_lemmas <$> hyps.mfilter is_acceptable_hyp

-- TODO mk_apps recursively
meta def inflate_under_apps (locals : list expr) : expr → tactic (list expr)
| e := do
  rws ← list.map prod.fst <$> mk_apps e locals,
  rws_extras ← list.join <$> rws.mmap inflate_under_apps,
  return $ e :: (rws ++ rws_extras)

meta def inflate_rw (locals : list expr) : expr × bool → tactic (list (expr × bool))
| (e, sy) := do
  as ← inflate_under_apps locals e,
  return $ as.map $ λ a, (a, sy)

meta def is_rewrite_lemma (d : declaration) : option (name × expr) :=
  let t := d.type in if is_acceptable_rewrite t then some (d.to_name, t) else none

meta def find_all_rewrites : tactic (list (name × expr)) := do
  e ← get_env,
  return $ e.decl_omap is_rewrite_lemma

end tidy.rewrite_search.discovery


import data.cpi.name
import data.list.witness

namespace cpi

/-- A telescope may extend a context by 0 or 1 levels. This is effectively a
    function on contexts, but having it be inductive allows us to case split on
    it, simplifying some proofs. -/
@[derive decidable_eq, nolint has_inhabited_instance]
inductive telescope : Type
| extend : ℕ → telescope
| preserve : telescope

/-- Apply a telescope to a context. -/
def telescope.apply : telescope → context → context
| (telescope.extend n) Γ := context.extend n Γ
| telescope.preserve Γ := Γ

/-- A prefix expression. This can either be one of:

  - A communication prefix (send a series of variables on a channel, and then
    recieve, binding $n$ variables).

  - A spontaneous or silent prefix: a spontaneous reaction with some rate $k$.
    Used to model when a molecule may decompose into a simpler one.

  The prefix is parameterised by two types: the context it exists in, a function
  which augments an arbitrary context with variables bound by this prefix.

  While it is possible to do without the second parameter, this introduces many
  complexities to the proof when renaming, as you need to
  `augment (rename π) = augment π', while preserving type safety.
-/
@[derive decidable_eq, nolint has_inhabited_instance]
inductive prefix_expr (ℍ : Type) : context → telescope → Type
| communicate {} {Γ} (a :  name Γ) (b : list (name Γ)) (y : ℕ) : prefix_expr Γ (telescope.extend y)
| spontaneous {} {Γ} (k : ℍ) : prefix_expr Γ telescope.preserve

variables {ℍ : Type}

-- Define some additional notation, and sugar
notation a `#(` b ` ; ` y `)` := prefix_expr.communicate a b y
notation a `#(` y `)` := prefix_expr.communicate a [] y
notation a `#⟨` b `⟩` := prefix_expr.communicate a b 0
notation a `#` := prefix_expr.communicate a [] 0

notation `τ@`:max k:max := prefix_expr.spontaneous k

namespace prefix_expr
  /-- Convert a prefix expression to a string. Can use `repr` normally. -/
  protected def to_string [has_repr ℍ] {Γ} : ∀ {f}, prefix_expr ℍ Γ f → string
  | ._ (a #( []; 0)) := repr a ++ "#"
  | ._ (a #( b; 0)) := repr a ++ "#<  " ++ repr b ++ "⟩"
  | ._ (a #( []; y)) := repr a ++ "#(" ++ repr y ++ ")"
  | ._ (a #( b; y)) := repr a ++ "#(" ++ repr b ++ ";" ++ repr y ++ ")"
  | ._ (τ@ k) := "τ@" ++ repr k

  instance [has_repr ℍ] {Γ f} : has_repr (prefix_expr ℍ Γ f) := ⟨ prefix_expr.to_string ⟩

  /-- A wrapper for prefixed expressions, which hides the extension function.

      This is suitable for comparing prefixes. -/
  @[nolint has_inhabited_instance]
  inductive wrap (ℍ : Type) : context → Type
  | intro {} {Γ} {f} (π : prefix_expr ℍ Γ f) : wrap Γ

  section free
    /-- Determine if any variable with a given level occurs within this prefix.
    -/
    def free_in : ∀ {Γ} {f}, level Γ → prefix_expr ℍ Γ f → Prop
    | ._ ._ l (a#(b; y)) := l ∈ a ∨ ∃ x ∈ b, l ∈ x
    | ._ ._ l τ@_ := false

    instance {Γ} {f} : has_mem (level Γ) (prefix_expr ℍ Γ f) := ⟨ free_in ⟩

    private def free_in_decide : ∀ {Γ} {f}
      (l : level Γ) (π : prefix_expr ℍ Γ f), decidable (free_in l π)
    | Γ ._ l (a#(b; y)) := if h : l ∈ a ∨ ∃ x ∈ b, l ∈ x then is_true h else is_false h
    | ._ ._ l τ@_ := decidable.false

    instance free_in.decidable {Γ} {f} {l} {π : prefix_expr ℍ Γ f}
      : decidable (free_in l π)
      := free_in_decide l π
  end free

  /- Renaming and extension.

      The extension sections are largely similar to that in data.cpi.name - see
      there for an explanation of some of the decisions made. -/
  section rename
    /-- Raise a level according to this prefix's context extension function. -/
    def raise :
      ∀ {Γ η} {f} (π : prefix_expr ℍ η f)
      , level Γ → level (f.apply Γ)
    | Γ ._ ._ (a#(b; y)) l := level.extend l
    | Γ ._ ._ τ@_ l := l

    /-- Rename all names within a prefix expression, providing some witness that
        this variable is free within it. -/
    def rename_with {Γ Δ} :
      Π {f} (π : prefix_expr ℍ Γ f)
      , (Π (a : name Γ), name.to_level a ∈ π → name Δ) → prefix_expr ℍ Δ f
    | f (a#(b; y)) ρ :=
      let a' := ρ a (or.inl (name.to_level_at a)) in
      let b' := list.map_witness b
        (λ x mem, ρ x (or.inr ⟨ x, mem, name.to_level_at x ⟩))
      in
      a'#( b' ; y)
    | f τ@k ρ := τ@k

    /-- Simple renaming function, not taking any witness. -/
    @[reducible]
    def rename {Γ Δ} {f} (ρ : name Γ → name Δ) (π : prefix_expr ℍ Γ f)
      : prefix_expr ℍ Δ f
      := rename_with π (λ a _, ρ a)

    /-- Renaming with the identity function does nothing. -/
    lemma rename_with_id {Γ} : ∀ {f} (π : prefix_expr ℍ Γ f)
      , rename_with π (λ a _, a) = π
    | ._ (a#(b; y)) := by simp [rename_with]
    | ._ τ@k := rfl

    /-- Renaming with the identity function is the identity. -/
    lemma rename_id {Γ} {f} (π : prefix_expr ℍ Γ f) : rename id π = π
      := rename_with_id π

    /-- Renaming twice is the same as renaming with a composed function. -/
    lemma rename_with_compose {Γ Δ η} :
      ∀ {f} (π : prefix_expr ℍ Γ f)
        (ρ : Π (a : name Γ), name.to_level a ∈ π → name Δ)
        (σ : name Δ → name η)
      , rename σ (rename_with π ρ) = rename_with π (λ a free, σ (ρ a free))
    | f (a#(b; y)) ρ σ := by simp [rename_with, rename, list.map_witness_to_map]
    | f (τ@_) ρ σ := rfl

    /-- Renaming twice is the same as renaming with a composed function. -/
    lemma rename_compose {Γ Δ η} {f}
        (π : prefix_expr ℍ Γ f) (ρ : name Γ → name Δ) (σ : name Δ → name η)
      : rename σ (rename ρ π) = rename (σ ∘ ρ) π
    := rename_with_compose π _ _

    /-- Wrap a renaming function, making it suitable for a nested context. -/
    def ext_with {Γ Δ η} :
      ∀ {f} (π : prefix_expr ℍ η f)
        (P : level (f.apply Γ) → Prop)
        (ρ : Π (x : name Γ), P (prefix_expr.raise π (name.to_level x)) → name Δ)
      , Π (x : name (f.apply Γ)), P (name.to_level x) → name (f.apply Δ)
    | f (_#(_; y)) P ρ a p := name.ext_with P ρ a p
    | f τ@_ P ρ a p := ρ a p

    /-- Extending with the identity does nothing. -/
    lemma ext_with_identity :
      ∀ {Γ η} {f} (π : prefix_expr ℍ η f)
        (P : level (f.apply Γ) → Prop)
        (a : name (f.apply Γ)) (p : P (name.to_level a))
      , ext_with π P (λ x _, x) a p = a
    | Γ η ._ (_#(_; _)) P a p := name.ext_with_identity P a p
    | Γ η ._ τ@k P a p := rfl

    /-- Extending with the identity does nothing. -/
    lemma ext_with_id {Γ η} {f}
        (π : prefix_expr ℍ η f) (P : level (f.apply Γ) → Prop)
      : ext_with π P (λ x _, x) = λ x _, x
      := funext $ λ a, funext (ext_with_identity π P a)

    /-- Wrap a simple renaming function, making it suitable for a nested context. -/
    def ext {Γ Δ η} {f} (π : prefix_expr ℍ η f) (ρ : name Γ → name Δ)
          : name (f.apply Γ) → name (f.apply Δ)
    | a := ext_with π (λ _, true) (λ x _, ρ x) a true.intro

    /-- Extending with the identity does nothing. -/
    lemma ext_identity {Γ η} {f} (π : prefix_expr ℍ η f) (a : name (f.apply Γ))
      : ext π id a = a := ext_with_identity π _ a _

    /-- Extending with the identity yields the identity function. -/
    lemma ext_id : ∀ {Γ η} {f} (π : prefix_expr ℍ η f), @ext ℍ Γ Γ η f π id = id
    | Γ η f π := funext (ext_identity π)

    /-- Composing extensions is equivalent extending a composition. -/
    lemma ext_with_compose :
      ∀ {Γ Δ η φ} {f} (π : prefix_expr ℍ φ f)
        (P : level (f.apply Γ) → Prop)
        (ρ : Π (x : name Γ), P (raise π (name.to_level x)) → name Δ)
        (σ : name Δ → name η)
        (a : name (f.apply Γ)) (p : P (name.to_level a))
      , ext π σ (ext_with π P ρ a p) = ext_with π P (λ a p, σ (ρ a p)) a p
    | Γ Δ η φ f (_#(_;_)) P ρ σ a p := name.ext_with_compose P ρ σ a p
    | Γ Δ η φ f τ@_ P ρ σ _ _ := rfl

    /-- Composing extensions is equivalent extending a composition. -/
    lemma ext_with_comp {Γ Δ η φ} {f} (π : prefix_expr ℍ φ f)
        (P : level (f.apply Γ) → Prop)
        (ρ : Π (x : name Γ), P (raise π (name.to_level x)) → name Δ)
        (σ : name Δ → name η)
      : (λ a p, ext π σ (ext_with π P ρ a p)) = ext_with π P (λ a p, σ (ρ a p))
      := funext $ λ a, funext (ext_with_compose π P ρ σ a)

    /-- Composing extensions is equivalent extending a composition. -/
    lemma ext_compose :
      ∀ {Γ Δ η φ} {f} (ρ : name Γ → name Δ) (σ : name Δ → name η)
        (π : prefix_expr ℍ φ f) (α : name (f.apply Γ))
      , ext π σ (ext π ρ α) = ext π (σ ∘ ρ) α
    | Γ Δ η φ f ρ σ (a#(b; y)) α := name.ext_compose ρ σ α
    | Γ Δ η φ f ρ σ τ@k α := rfl

    /-- Composing extensions is equivalent extending a composition. -/
    lemma ext_comp :
      ∀ {Γ Δ η φ} {f} (ρ : name Γ → name Δ) (σ : name Δ → name η)
        (π : prefix_expr ℍ φ f)
      , (ext π σ ∘ ext π ρ) = ext π (σ ∘ ρ)
    | Γ Δ η φ f ρ σ π := funext (ext_compose ρ σ π)

    /-- Rewrite one ext_with to another -/
    lemma ext_with_discard :
      ∀ {Γ Δ η} {f} (π : prefix_expr ℍ η f)
        (P : level (f.apply Γ) → Prop)
        (ρ : name Γ → name Δ)
      , (ext_with π P (λ a _, ρ a))
      = (λ a _, ext_with π (λ _x, true) (λ x _, ρ x) a true.intro)
    | Γ Δ η f (a'#(b; y)) P ρ := funext $ λ a, funext $ λ free, begin
        have : (λ (a : name Γ) (_x : P (@raise ℍ _ _ _ (a'#(b ; y)) (name.to_level a))), ρ a)
             = (λ (a : name Γ) (_x : P (level.extend (name.to_level a))), ρ a)
             := rfl,
        unfold ext_with,
        rw [this, name.ext_with_discard P],
        from rfl,
      end
    | Γ Δ η f τ@_ P ρ := funext $ λ a, funext $ λ free, rfl

    /-- Raising with a renamed prefix has the same effect as the original one. -/
    lemma rename_with_raise
        {Γ Δ η} {f} (π : prefix_expr ℍ Γ f)
        (ρ : Π (x : name Γ), name.to_level x ∈ π → name Δ)
        (l : level η)
      : raise π l = raise (rename_with π ρ) l
      := by { cases π; from rfl }

    /-- Extending with a renamed prefix has the same effect as the original one. -/
    lemma rename_with_ext_with
        {Γ Δ η φ} {f} (π : prefix_expr ℍ η f)
        (P : level (f.apply Γ) → Prop)
        (ρ : name Γ → name Δ)
        (σ : Π (x : name η), name.to_level x ∈ π → name φ)
      : ext_with (rename_with π σ) P (λ a _, ρ a) = ext_with π P (λ a _, ρ a)
      := funext $ λ a, funext $ λ free, by { cases π; from rfl }

    /-- Extending with a renamed prefix has the same effect as the original one. -/
    lemma rename_ext
        {Γ Δ η φ} {f} (ρ : name Γ → name Δ) (σ : name η → name φ)
        (π : prefix_expr ℍ Γ f)
      : @ext ℍ η φ Γ f π σ = (ext (rename ρ π) σ)
      := funext $ λ a, by { cases π; from rfl }
  end rename

  section rename_equations
    variables {Γ Δ : context} {ρ : name Γ → name Δ}

    @[simp]
    lemma rename_communicate (a :  name Γ) (b : list (name Γ)) (y : ℕ)
      : rename ρ (a#(b; y) : prefix_expr ℍ _ _) = ((ρ a)#(list.map ρ b; y))
      := by simp [rename, rename_with, list.map_witness_to_map]

    @[simp]
    lemma rename_spontaneous (k : ℍ)
      : rename ρ τ@k = τ@k
      := by simp only [rename, rename_with]

    @[simp]
    lemma ext_communicate {η} (a :  name η) (b : list (name η)) (y : ℕ)
      : ext (a#(b; y) : prefix_expr ℍ _ _) ρ = name.ext ρ
      := funext $ λ x, by { unfold ext ext_with, from rfl }

    @[simp]
    lemma ext_spontaneous {η} (k : ℍ)
      : ext (@spontaneous ℍ η k) ρ = ρ
      := funext $ λ x, by unfold ext ext_with

    private lemma rename_inj {Γ Δ} {ρ : name Γ → name Δ} (inj : function.injective ρ)
      : ∀ {f₁ f₂} {π₁ : prefix_expr ℍ Γ f₁} {π₂ : prefix_expr ℍ Γ f₂}
      , wrap.intro (rename ρ π₁) = wrap.intro (rename ρ π₂)
      → π₁ == π₂
    | _ _ (a#(b; y)) (a'#(b'; y')) eq := begin
        simp only [rename_communicate] at eq,
        rcases eq with ⟨ ⟨ _ ⟩, eqπ ⟩,
        simp only [heq_iff_eq] at eqπ ⊢,
        from ⟨ inj eqπ.left, list.injective_map_iff.mpr inj eqπ.right ⟩,
      end
    | _ _ (a#(b; y)) τ@k eq := begin
      simp only [rename_communicate, rename_spontaneous] at eq,
      exfalso, from eq.1,
    end
    | _ _ τ@k (a#(b; y)) eq := begin
        simp only [rename_communicate, rename_spontaneous] at eq,
        exfalso, from eq.1,
      end
    | _ _ τ@k τ@k' eq := begin
        simp only [rename_spontaneous] at eq,
        rcases eq with ⟨ eqC, eqπ ⟩,
        cases eqπ, from heq.refl _,
      end

    lemma rename.inj {Γ Δ} {ρ : name Γ → name Δ} (inj : function.injective ρ)
      : ∀ {f}, function.injective (@rename ℍ Γ Δ f ρ)
    | f π₁ π₂ eq := eq_of_heq (rename_inj inj (congr_arg wrap.intro eq))

    lemma ext.inj {Γ Δ η} {ρ : name Γ → name Δ}
      : ∀ {f} (π : prefix_expr ℍ η f), function.injective ρ → function.injective (ext π ρ)
    | ._ (_#(_; y)) inj a b eq := begin
        simp only [ext_communicate] at eq,
        from name.ext.inj inj eq,
      end
    | ._ (τ@_) inj a b eq := by { simp only [ext_spontaneous] at eq, from inj eq, }
  end rename_equations
end prefix_expr

end cpi

#lint-


/-
Copyright (c) 2016 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
prelude
import init.control
import init.meta.simp_tactic
import init.meta.smt.congruence_closure
import init.meta.smt.ematch

universe u

run_cmd mk_simp_attr `pre_smt
run_cmd mk_hinst_lemma_attr_set `ematch [] [`ematch_lhs]

/--
  Configuration for the smt tactic preprocessor. The preprocessor
  is applied whenever a new hypothesis is introduced.

  - simp_attr: is the attribute name for the simplification lemmas
    that are used during the preprocessing step.

  - max_steps: it is the maximum number of steps performed by the simplifier.

  - zeta: if tt, then zeta reduction (i.e., unfolding let-expressions)
    is used during preprocessing.
-/
structure smt_pre_config :=
(simp_attr : name := `pre_smt)
(max_steps : nat  := 1000000)
(zeta      : bool := ff)

/--
Configuration for the smt_state object.

- em_attr: is the attribute name for the hinst_lemmas
  that are used for ematching -/
structure smt_config :=
(cc_cfg        : cc_config      := {})
(em_cfg        : ematch_config  := {})
(pre_cfg       : smt_pre_config := {})
(em_attr       : name           := `ematch)

meta def smt_config.set_classical (c : smt_config) (b : bool) : smt_config :=
{ cc_cfg := { em := b, ..c.cc_cfg }, ..c }

meta constant smt_goal                  : Type
meta def smt_state :=
list smt_goal
meta constant smt_state.mk              : smt_config → tactic smt_state
meta constant smt_state.to_format       : smt_state → tactic_state → format
/-- Return tt iff classical excluded middle was enabled at  smt_state.mk -/
meta constant smt_state.classical       : smt_state → bool

meta def smt_tactic :=
state_t smt_state tactic

meta instance : has_append smt_state :=
list.has_append

section
local attribute [reducible] smt_tactic
meta instance : monad smt_tactic := by apply_instance
meta instance : alternative smt_tactic := by apply_instance
meta instance : monad_state smt_state smt_tactic := by apply_instance
end

/- We don't use the default state_t lift operation because only
   tactics that do not change hypotheses can be automatically lifted to smt_tactic. -/
meta constant tactic_to_smt_tactic (α : Type) : tactic α → smt_tactic α

meta instance : has_monad_lift tactic smt_tactic :=
⟨tactic_to_smt_tactic⟩

meta instance (α : Type) : has_coe (tactic α) (smt_tactic α) :=
⟨monad_lift⟩

meta instance : monad_fail smt_tactic :=
{ fail := λ α s, (tactic.fail (to_fmt s) : smt_tactic α), ..smt_tactic.monad }

namespace smt_tactic
open tactic (transparency)
meta constant intros                     : smt_tactic unit
meta constant intron                     : nat  → smt_tactic unit
meta constant intro_lst                  : list name → smt_tactic unit
/--
  Try to close main goal by using equalities implied by the congruence
  closure module.
-/
meta constant close                           : smt_tactic unit
/--
  Produce new facts using heuristic lemma instantiation based on E-matching.
  This tactic tries to match patterns from lemmas in the main goal with terms
  in the main goal. The set of lemmas is populated with theorems
  tagged with the attribute specified at smt_config.em_attr, and lemmas
  added using tactics such as `smt_tactic.add_lemmas`.
  The current set of lemmas can be retrieved using the tactic `smt_tactic.get_lemmas`.

  Remark: the given predicate is applied to every new instance. The instance
  is only added to the state if the predicate returns tt.
-/
meta constant ematch_core                     : (expr → bool) → smt_tactic unit
/--
  Produce new facts using heuristic lemma instantiation based on E-matching.
  This tactic tries to match patterns from the given lemmas with terms in
  the main goal.
-/
meta constant ematch_using                    : hinst_lemmas → smt_tactic unit
meta constant mk_ematch_eqn_lemmas_for_core   : transparency → name → smt_tactic hinst_lemmas
meta constant to_cc_state                     : smt_tactic cc_state
meta constant to_em_state                     : smt_tactic ematch_state
meta constant get_config                      : smt_tactic smt_config
/--
  Preprocess the given term using the same simplifications rules used when
  we introduce a new hypothesis. The result is pair containing the resulting
  term and a proof that it is equal to the given one.
-/
meta constant preprocess                      : expr → smt_tactic (expr × expr)
meta constant get_lemmas                      : smt_tactic hinst_lemmas
meta constant set_lemmas                      : hinst_lemmas → smt_tactic unit
meta constant add_lemmas                      : hinst_lemmas → smt_tactic unit

meta def add_ematch_lemma_core (md : transparency) (as_simp : bool) (e : expr) : smt_tactic unit :=
do h  ← hinst_lemma.mk_core md e as_simp,
   add_lemmas (mk_hinst_singleton h)

meta def add_ematch_lemma_from_decl_core (md : transparency) (as_simp : bool) (n : name) : smt_tactic unit :=
do h  ← hinst_lemma.mk_from_decl_core md n as_simp,
   add_lemmas (mk_hinst_singleton h)

meta def add_ematch_eqn_lemmas_for_core  (md : transparency) (n : name) : smt_tactic unit :=
do hs ← mk_ematch_eqn_lemmas_for_core md n,
   add_lemmas hs

meta def ematch : smt_tactic unit :=
ematch_core (λ _, tt)

meta def failed {α} : smt_tactic α :=
tactic.failed

meta def fail {α : Type} {β : Type u} [has_to_format β] (msg : β) : smt_tactic α :=
tactic.fail msg

meta def try {α : Type} (t : smt_tactic α) : smt_tactic unit :=
⟨λ ss ts, result.cases_on (t.run ss ts)
 (λ ⟨a, new_ss⟩, result.success ((), new_ss))
 (λ e ref s', result.success ((), ss) ts)⟩

/-- `iterate_at_most n t`: repeat the given tactic at most n times or until t fails -/
meta def iterate_at_most : nat → smt_tactic unit → smt_tactic unit
| 0     t := return ()
| (n+1) t := (do t, iterate_at_most n t) <|> return ()

/-- `iterate_exactly n t` : execute t n times -/
meta def iterate_exactly : nat → smt_tactic unit → smt_tactic unit
| 0     t := return ()
| (n+1) t := do t, iterate_exactly n t

meta def iterate : smt_tactic unit → smt_tactic unit :=
iterate_at_most 100000

meta def eblast : smt_tactic unit :=
iterate (ematch >> try close)

open tactic

protected meta def read : smt_tactic (smt_state × tactic_state) :=
do s₁ ← get,
   s₂ ← tactic.read,
   return (s₁, s₂)

protected meta def write : smt_state × tactic_state → smt_tactic unit :=
λ ⟨ss, ts⟩, ⟨λ _ _, result.success ((), ss) ts⟩

private meta def mk_smt_goals_for (cfg : smt_config) : list expr → list smt_goal → list expr
                                  → tactic (list smt_goal × list expr)
| []        sr tr := return (sr.reverse, tr.reverse)
| (tg::tgs) sr tr := do
  tactic.set_goals [tg],
  [new_sg] ← smt_state.mk cfg | tactic.failed,
  [new_tg] ← get_goals | tactic.failed,
  mk_smt_goals_for tgs (new_sg::sr) (new_tg::tr)

/- See slift -/
meta def slift_aux {α : Type} (t : tactic α) (cfg : smt_config) : smt_tactic α :=
⟨λ ss, do
   _::sgs  ← return ss | tactic.fail "slift tactic failed, there no smt goals to be solved",
   tg::tgs ← tactic.get_goals | tactic.failed,
   tactic.set_goals [tg], a ← t,
   new_tgs ← tactic.get_goals,
   (new_sgs, new_tgs) ← mk_smt_goals_for cfg new_tgs [] [],
   tactic.set_goals (new_tgs ++ tgs),
   return (a, new_sgs ++ sgs)⟩

/--
  This lift operation will restart the SMT state.
  It is useful for using tactics that change the set of hypotheses. -/
meta def slift {α : Type} (t : tactic α) : smt_tactic α :=
get_config >>= slift_aux t

meta def trace_state : smt_tactic unit :=
do (s₁, s₂) ← smt_tactic.read,
   trace (smt_state.to_format s₁ s₂)

meta def trace {α : Type} [has_to_tactic_format α] (a : α) : smt_tactic unit :=
tactic.trace a

meta def to_expr (q : pexpr) (allow_mvars := tt) : smt_tactic expr :=
tactic.to_expr q allow_mvars

meta def classical : smt_tactic bool :=
do s ← get,
   return s.classical

meta def num_goals : smt_tactic nat :=
list.length <$> get

/- Low level primitives for managing set of goals -/
meta def get_goals : smt_tactic (list smt_goal × list expr) :=
do (g₁, _) ← smt_tactic.read,
   g₂ ← tactic.get_goals,
   return (g₁, g₂)

meta def set_goals : list smt_goal → list expr → smt_tactic unit :=
λ g₁ g₂, ⟨λ ss, tactic.set_goals g₂ >> return ((), g₁)⟩

private meta def all_goals_core (tac : smt_tactic unit) : list smt_goal → list expr → list smt_goal → list expr → smt_tactic unit
| []        ts        acs act := set_goals acs (ts ++ act)
| (s :: ss) []        acs act := fail "ill-formed smt_state"
| (s :: ss) (t :: ts) acs act :=
  do set_goals [s] [t],
     tac,
     (new_ss, new_ts) ← get_goals,
     all_goals_core ss ts (acs ++ new_ss) (act ++ new_ts)

/-- Apply the given tactic to all goals. -/
meta def all_goals (tac : smt_tactic unit) : smt_tactic unit :=
do (ss, ts) ← get_goals,
   all_goals_core tac ss ts [] []

/-- LCF-style AND_THEN tactic. It applies tac1, and if succeed applies tac2 to each subgoal produced by tac1 -/
meta def seq (tac1 : smt_tactic unit) (tac2 : smt_tactic unit) : smt_tactic unit :=
do (s::ss, t::ts) ← get_goals,
   set_goals [s] [t],
   tac1, all_goals tac2,
   (new_ss, new_ts) ← get_goals,
   set_goals (new_ss ++ ss) (new_ts ++ ts)

meta instance : has_andthen (smt_tactic unit) (smt_tactic unit) (smt_tactic unit) :=
⟨seq⟩

meta def focus1 {α} (tac : smt_tactic α) : smt_tactic α :=
do (s::ss, t::ts) ← get_goals,
   match ss with
   | []  := tac
   | _   := do
     set_goals [s] [t],
     a ← tac,
     (ss', ts') ← get_goals,
     set_goals (ss' ++ ss) (ts' ++ ts),
     return a
   end

meta def solve1 (tac : smt_tactic unit) : smt_tactic unit :=
do (ss, gs) ← get_goals,
   match ss, gs with
   | [],     _    := fail "solve1 tactic failed, there isn't any goal left to focus"
   | _,     []    := fail "solve1 tactic failed, there isn't any smt goal left to focus"
   | s::ss, g::gs :=
     do set_goals [s] [g],
        tac,
        (ss', gs') ← get_goals,
        match ss', gs' with
        | [], [] := set_goals ss gs
        | _,  _  := fail "solve1 tactic failed, focused goal has not been solved"
        end
   end

meta def swap : smt_tactic unit :=
do (ss, ts) ← get_goals,
   match ss, ts with
   | (s₁ :: s₂ :: ss), (t₁ :: t₂ :: ts) := set_goals (s₂ :: s₁ :: ss) (t₂ :: t₁ :: ts)
   | _,                _                := failed
   end

/-- Add a new goal for t, and the hypothesis (h : t) in the current goal. -/
meta def assert (h : name) (t : expr) : smt_tactic unit :=
tactic.assert_core h t >> swap >> intros >> swap >> try close

/-- Add the hypothesis (h : t) in the current goal if v has type t. -/
meta def assertv (h : name) (t : expr) (v : expr) : smt_tactic unit :=
tactic.assertv_core h t v >> intros >> return ()

/-- Add a new goal for t, and the hypothesis (h : t := ?M) in the current goal. -/
meta def define  (h : name) (t : expr) : smt_tactic unit :=
tactic.define_core h t >> swap >> intros >> swap >> try close

/-- Add the hypothesis (h : t := v) in the current goal if v has type t. -/
meta def definev (h : name) (t : expr) (v : expr) : smt_tactic unit :=
tactic.definev_core h t v >> intros >> return ()

/-- Add (h : t := pr) to the current goal -/
meta def pose (h : name) (t : option expr := none) (pr : expr) : smt_tactic unit :=
match t with
| none   := do t ← infer_type pr, definev h t pr
| some t := definev h t pr
end

/-- Add (h : t) to the current goal, given a proof (pr : t) -/
meta def note (h : name) (t : option expr := none) (pr : expr) : smt_tactic unit :=
match t with
| none   := do t ← infer_type pr, assertv h t pr
| some t := assertv h t pr
end

meta def destruct (e : expr) : smt_tactic unit :=
smt_tactic.seq (tactic.destruct e) smt_tactic.intros

meta def by_cases (e : expr) : smt_tactic unit :=
do c ← classical,
   if c then
     destruct (expr.app (expr.const `classical.em []) e)
   else do
     dec_e ← (mk_app `decidable [e] <|> fail "by_cases smt_tactic failed, type is not a proposition"),
     inst  ← (mk_instance dec_e <|> fail "by_cases smt_tactic failed, type of given expression is not decidable"),
     em    ← mk_app `decidable.em [e, inst],
     destruct em

meta def by_contradiction : smt_tactic unit :=
do t ← target,
   c ← classical,
   if t.is_false then skip
   else if c then do
      apply (expr.app (expr.const `classical.by_contradiction []) t),
      intros
   else do
     dec_t ← (mk_app `decidable [t] <|> fail "by_contradiction smt_tactic failed, target is not a proposition"),
     inst  ← (mk_instance dec_t <|> fail "by_contradiction smt_tactic failed, target is not decidable"),
     a     ← mk_mapp `decidable.by_contradiction [some t, some inst],
     apply a,
     intros

/-- Return a proof for e, if 'e' is a known fact in the main goal. -/
meta def proof_for (e : expr) : smt_tactic expr :=
do cc ← to_cc_state, cc.proof_for e

/-- Return a refutation for e (i.e., a proof for (not e)), if 'e' has been refuted in the main goal. -/
meta def refutation_for (e : expr) : smt_tactic expr :=
do cc ← to_cc_state, cc.refutation_for e

meta def get_facts : smt_tactic (list expr) :=
do cc ← to_cc_state,
   return $ cc.eqc_of expr.mk_true

meta def get_refuted_facts : smt_tactic (list expr) :=
do cc ← to_cc_state,
   return $ cc.eqc_of expr.mk_false

meta def add_ematch_lemma : expr → smt_tactic unit :=
add_ematch_lemma_core reducible ff

meta def add_ematch_lhs_lemma : expr → smt_tactic unit :=
add_ematch_lemma_core reducible tt

meta def add_ematch_lemma_from_decl : name → smt_tactic unit :=
add_ematch_lemma_from_decl_core reducible ff

meta def add_ematch_lhs_lemma_from_decl : name → smt_tactic unit :=
add_ematch_lemma_from_decl_core reducible ff

meta def add_ematch_eqn_lemmas_for : name → smt_tactic unit :=
add_ematch_eqn_lemmas_for_core reducible

meta def add_lemmas_from_facts_core : list expr → smt_tactic unit
| []      := return ()
| (f::fs) := do
  try (is_prop f >> guard (f.is_pi && bnot (f.is_arrow)) >> proof_for f >>= add_ematch_lemma_core reducible ff),
  add_lemmas_from_facts_core fs

meta def add_lemmas_from_facts : smt_tactic unit :=
get_facts >>= add_lemmas_from_facts_core

meta def induction (e : expr) (ids : list name := []) (rec : option name := none) : smt_tactic unit :=
slift (tactic.induction e ids rec >> return ()) -- pass on the information?

meta def when (c : Prop) [decidable c] (tac : smt_tactic unit) : smt_tactic unit :=
if c then tac else skip

meta def when_tracing (n : name) (tac : smt_tactic unit) : smt_tactic unit :=
when (is_trace_enabled_for n = tt) tac

end smt_tactic

open smt_tactic

meta def using_smt {α} (t : smt_tactic α) (cfg : smt_config := {}) : tactic α :=
do ss ← smt_state.mk cfg,
   (a, _) ← (do a ← t, iterate close, return a).run ss,
   return a

meta def using_smt_with {α} (cfg : smt_config) (t : smt_tactic α) : tactic α :=
using_smt t cfg