import Mathlib.Tactic.Linarith
import Aesop
-- A translation from SSA + Regions to Tree, with no proofs.
-- this allows for easy unfolding of the semantics.
-- If we carry around proofs of well formedness, the dependent typing
-- of the well-formedness leads to stuck terms.
-- Thus, we eschew the need to have proofs, and simply YOLO translate from
-- the origina SSA + Regions program into Tree.

namespace SSARgnVar2Tree

inductive RgnName
| r0
| r1
| r2
| r3
| r4
deriving DecidableEq

inductive VarName
| null
| x0
| x1
| x2
| x3
| x4
| x5
| x6
| x7
| y1
| y2
| y3
| y4
| z1
| z2
| z3
| z4
deriving DecidableEq

section StxSem

-- | The environment is a mapping from variable names to values.
abbrev Env (k : Type) (α : Type) := k → α

-- | Extend the environment with a new variable.
@[simp]
def Env.extend [DecidableEq k] (name : k) (v : α) (e: Env k α) : Env k α :=
  fun name' => if name = name' then v else e name'

@[simp]
def Env.empty [Inhabited α] : Env k α :=
  fun _ => default

@[simp]
def Env.map (f : α → β) (e : Env k α) : Env k β :=
  fun name => f (e name)

notation "∅" => Env.empty
notation e "[" name "↦" v "]" => Env.extend name v e

abbrev VarEnv (α : Type) := Env VarName α
abbrev RgnEnv (α : Type) := Env RgnName α

-- The input semantics given by the user.
class UserSemantics (opcode : Type)  where
  -- | Arguments given as (<args>, <rgn1>, ... <rgnN>)
  -- | Consider not allowing users to not be access all of 'Env'.
  opcodeEval (op: opcode) (vals : Int × Int) (rgns : (Int × Int → Int)) : Int

attribute [simp] UserSemantics.opcodeEval

inductive ASTKind : Type where
| O : ASTKind
| Os : ASTKind
| R : ASTKind
deriving Inhabited, DecidableEq

-- | The operations of the language.
inductive AST (opcode : Type): ASTKind → Type where
| assign (ret : VarName) (op : opcode)
  (args : VarName × VarName)
  (rgns : AST opcode .R) : AST opcode .O
| ops1 (op : AST opcode .O) :  AST opcode .Os
| opsmany (op : AST opcode .O) (ops : AST opcode .Os) : AST opcode .Os
| rgn (args : VarName × VarName) (body : AST opcode .Os) : AST opcode .R
| rgnvar (var : RgnName) : AST opcode .R
| rgn0 : AST opcode .R

def AST.retname : AST opcode .O → VarName
| .assign (ret := ret) .. => ret

instance [Inhabited opcode] : Inhabited (AST opcode .O) where
  default := .assign .x0 default (.x0, .x0) .rgn0

instance [Inhabited opcode] : Inhabited (AST opcode .Os) where
  default := .ops1 default

instance [Inhabited opcode] : Inhabited (AST opcode .R) where
  default := .rgn (.x0, .x0) default

@[simp]
def Ops.ofList [Inhabited opcode] : List (AST opcode .O) → AST opcode .Os
| [] => panic! "need non-empty list"
| [x] => .ops1 x
| x :: xs => .opsmany x (Ops.ofList xs)

@[reducible, simp]
instance [Inhabited opcode] : Coe (List (AST opcode .O)) (AST opcode .Os) := ⟨Ops.ofList⟩

@[simp, reducible]
def ASTKind.eval : ASTKind → Type
| .O => Int
| .Os => Int
| .R => (Int × Int → Int)
-- evaluate an operation with repect to a particular user semantics.
@[simp]
def AST.eval [S: UserSemantics opcode]
  {astk: ASTKind} (e: VarEnv Int)
  (re: RgnEnv (Int × Int → Int)): AST opcode astk → astk.eval × VarEnv Int
| .assign ret op args r =>
    let (arg1, arg2) := args
    let retval := S.opcodeEval op (e arg1, e arg2) (r.eval e re).fst
    (retval, e.extend ret retval)
| .ops1 op =>
   let (out, env) := op.eval e re
   (out, env)
| .opsmany op ops =>
    let e' := (op.eval e re).snd
    ops.eval e' re
| .rgnvar v => (re v, e)
| .rgn args body =>
    (fun vals =>
      let e := Env.empty
      let e1 := e.extend args.fst vals.fst
      let e2 := e1.extend args.snd vals.snd
      let (outval, _e) := body.eval e2 re
      outval, e)
| .rgn0 => (fun _ => default, e)

end StxSem

section Tree

inductive CtreeKind
| O -- op
| R -- region (higher order)
deriving Inhabited

-- closed trees, leaves are integers.
inductive Ctree  (opcode : Type) : CtreeKind → Type where
| binop
  (op : opcode)
  (lhs : Ctree opcode .O)
  (rhs : Ctree opcode .O)
  (rgns: Ctree opcode .R) : Ctree opcode .O
| rgn (f : Int × Int → Ctree opcode .O) : Ctree opcode .R
| rgn0 : Ctree opcode .R
| leaf (val : Int) : Ctree opcode .O

instance : Inhabited (Ctree opcode .O) where
  default := .leaf 10

instance : Inhabited (Ctree opcode .R) where
  default := .rgn0

-- convert an AST into a closed tree under the given environment
-- note that translation into a closed tree needs an environment,
-- to learn the values of variables.
-- This version has an `_` in the name since it needs a `Ctree.VarEnv`, not an
-- `Env`. We will writea helper that converts `Env` into `Ctree.VarEnv`.
@[simp, reducible]
def ASTKind.toCTree (opcode : Type) : ASTKind → Type
| .O => Ctree opcode .O
| .Os => Ctree opcode .O
| .R => Ctree opcode .R

@[simp]
def AST.toCtree_ {astk: ASTKind} (e : VarEnv (Ctree opcode .O))
  (re: RgnEnv (Ctree opcode .R)):
  AST opcode astk → (astk.toCTree opcode) × VarEnv (Ctree opcode .O)
| .assign ret opcode (u, v) r =>
    let rval := (r.toCtree_ e re).fst
    let t := .binop opcode (e u) (e v) rval
    (t, e.extend ret t)
| .rgnvar var => let e := Env.empty; (re var, e)
| .rgn0 => let e := Env.empty; (.rgn0, e)
| .rgn args body =>
    let e := Env.empty -- NOTE: regions are now isolated from above
    (.rgn fun vals =>
      let e1 := e.extend args.fst (.leaf vals.fst)
      let e2 := e1.extend args.snd (.leaf vals.snd)
      (body.toCtree_ e2 re).fst, e)
| .ops1 op =>
  let (val, e) := op.toCtree_ e re
  (val, e)
| .opsmany os o =>
    let (_, e) := os.toCtree_ e re
    o.toCtree_ e re

-- wrap every element in a (.leaf) constructor
@[simp]
def Ctree.VarEnv.ofEnv (e: VarEnv Int) : VarEnv (Ctree opcode .O) :=
   fun name => .leaf (e name)

-- TODO: should these be coercions?
@[simp]
def Ctree.RgnEnv.ofEnv (re: RgnEnv (Int × Int → Int)) :
  RgnEnv (Ctree opcode .R) :=
   fun name => Ctree.rgn (fun args => .leaf (re name args))

-- this converts an AST into a Ctree, given an environment
-- and an AST.
@[simp]
def Op.toCtree (a: AST opcode .O) (e: VarEnv Int)
  (re: RgnEnv (Int × Int → Int)) : Ctree opcode .O :=
    (a.toCtree_ (Ctree.VarEnv.ofEnv e) (Ctree.RgnEnv.ofEnv re)).fst

@[simp]
def Ops.toCtree (a: AST opcode .Os)
  (e: VarEnv Int) (re: RgnEnv (Int × Int → Int)) : Ctree opcode .O :=
    (a.toCtree_ (Ctree.VarEnv.ofEnv e) (Ctree.RgnEnv.ofEnv re)).fst

@[simp]
def Region.toCtree (a: AST opcode .R)
  (e: VarEnv Int)
  (re: RgnEnv (Int × Int → Int)): Ctree opcode .R :=
    (a.toCtree_ (Ctree.VarEnv.ofEnv e) (Ctree.RgnEnv.ofEnv re)).fst

-- evaluate a Ctree. note that this needs no environment.
@[simp]
def CtreeKind.eval : CtreeKind → Type
| .O => Int
| .R => Int × Int → Int

-- Note: is this literally "just" staging the partial evaluation against the environment?
def Ctree.eval [UserSemantics opcode] : Ctree opcode treek → treek.eval
| .binop o l r rs =>
  UserSemantics.opcodeEval o (l.eval, r.eval) rs.eval
| .leaf v => v
| .rgn0 => fun _ => default
| .rgn f => fun args => (f args).eval

end Tree

namespace MultipleInstructionTree
inductive Opcode
| add
| mul
| loop : Opcode
| ite : Opcode
| run : Opcode
| runnot : Opcode
| not : Opcode
| const : Int → Opcode
deriving Inhabited

def loopSemantics (n : Nat) (f : Int → Int) (v : Int) : Int :=
  match n with
  | 0 => v -- inline id
  | n + 1 => f (loopSemantics n f v) -- inline ∘

@[simp]
instance :  UserSemantics Opcode  where
  opcodeEval
  | .not, ⟨a, _⟩, _ => if a = 0 then 1 else 0
  | .mul, ⟨a, b⟩, _ => a * b
  | .add, ⟨a, b⟩, _ => a + b
  | .const i, ⟨_a, _b⟩, _ => i
  | .run, ⟨v, w⟩, r => r ⟨v, w⟩
  | .runnot, ⟨v, w⟩, r => r ⟨if v = 0 then 1 else 0, w⟩ -- execute: 'r(not v, w)'
  | .loop, ⟨n, init⟩, r => loopSemantics n.toNat (fun i => r ⟨i, 0⟩) init
  | _, _, _ => 42

def x_add_4_times_mul_val_eq (env: VarEnv Int):
  let p : AST Opcode .Os :=
      Ops.ofList [
      .assign .x1 .add (.x0, .x0) .rgn0,
      .assign .x2 .add (.x1, .x1) .rgn0
      ]
  let q : AST Opcode .Os := Ops.ofList [
        .assign .x1 (.const 4) (.x0, .x0) .rgn0
      , .assign .x2 .mul (.x1, .x0) .rgn0
    ]
  (Ops.toCtree p env renv).eval = (Ops.toCtree q env renv).eval := by {
    simp only [Ops.ofList, AST.eval, Ops.toCtree,  AST.toCtree_];
    -- see that there are environments, which are folded away when calling
    -- Ctree.eval.
    simp[Ctree.eval];
    linarith
  }

def run_inline :
  let p : AST Opcode .R :=
    AST.rgn ⟨.x0, .x1⟩ $ Ops.ofList [
      .assign .x2 .add (.x0, .x1) .rgn0 -- x2 := x0 + x1
      ]
  let q : AST Opcode .R :=
    AST.rgn ⟨.x0, .x1⟩ $ Ops.ofList [
      .assign .x2 .run (.x0, .x1) (AST.rgn ⟨.x5, .x6⟩ (Ops.ofList [
        -- x2 := run (x0, x1) { ^(x5, x6): return x5 + x6 }
        .assign .y1 .add (.x5, .x6) .rgn0
      ]))
    ]
  (Region.toCtree p Env.empty Env.empty).eval =
  (Region.toCtree q Env.empty Env.empty).eval := by {
    simp
    dsimp only[Ctree.eval]
  }

-- trying to convert an AST to a Ctree at any environment
-- is equivalent to converting it in an empty environment.
@[simp]
theorem AST.toCtree_rgn_equiv_empty
  (r: AST opcode .R)
  [Inhabited opcode] :
  AST.toCtree_ env renv r = AST.toCtree_ Env.empty renv r := by {
  simp; cases r <;> simp;
}

def runnot_inline:
  let p : AST Opcode .R :=
    AST.rgn ⟨.x0, .x1⟩ $ Ops.ofList [
      .assign .y1 .run ⟨.x0, .x1⟩ (.rgnvar .r1) -- y1 := r(x0, x1)
    ]
  (Region.toCtree p Env.empty Env.empty).eval =
  (Region.toCtree (.rgnvar .r1 : AST Opcode .R) Env.empty Env.empty).eval := by {
    simp
    simp[Ctree.eval];
    -- this stil cannot be simplified away, since this is in fact false.
    -- We need some notion of the fact that z1, z2 is free in 'r',
    -- so that we can say that we can remove 'z2, z1' in the 'q'
    -- environment when we run 'r'.
    -- if we have this, then we are good, and we can prove the theorem.
    -- but how do we convince people that use MLIR that this is a sensible
    -- thing to reason about?

  }

end MultipleInstructionTree

end SSARgnVar2Tree


import .fol .abel

universe u

-- section weekdays
-- @[derive has_reflect]
-- inductive weekday : Type
-- | monday          : weekday
-- | another_day     : weekday → weekday

-- open weekday

-- meta def dump_weekday (f : weekday) : tactic unit :=
-- tactic.trace $ to_string (expr.to_raw_fmt (reflect f).to_expr)

-- -- run_cmd dump_weekday (another_day (another_day monday))
-- --(app (const weekday.another_day []) (app (const weekday.another_day []) (const weekday.monday [])))

-- inductive weekday' : Type
-- | monday          : weekday'
-- | another_day     : weekday' → weekday'

-- open weekday'
-- meta instance has_reflect_weekday' : has_reflect weekday'
-- | weekday'.monday                  := `(monday)
-- | (weekday'.another_day x)         := `(λ l, weekday'.another_day l).subst $
--                                         by haveI := has_reflect_weekday'; exact (reflect x)

-- meta def dump_weekday' (f : weekday') : tactic unit :=
-- tactic.trace $ to_string (expr.to_raw_fmt (reflect f).to_expr)

-- -- run_cmd dump_weekday' (another_day (another_day monday))
-- -- (app (const weekday'.another_day []) (app (const weekday'.another_day []) (const weekday'.monday [])))
-- end weekdays
-- -- meta instance has_reflect_preterm {L : Language.{u}} : Π{n : ℕ}, has_reflect (preterm L n)
-- -- | 0 (var k) := `(@preterm.var L).subst (reflect k)

-- -- @[derive has_reflect]

-- open fol abel

-- section preterm_aux
-- inductive preterm_aux (L : Language.{u}) : Type u
-- | var     : ℕ → preterm_aux
-- | func    : ∀ k : ℕ, L.functions k → preterm_aux
-- | app     : preterm_aux → preterm_aux → preterm_aux

-- def to_aux {L : Language.{u}} : ∀ {l : ℕ},  preterm L l → preterm_aux L
-- | 0 (var n)      := preterm_aux.var _ n
-- | k (func f)     := preterm_aux.func _ f
-- | k (app t₁ t₂)  := preterm_aux.app (to_aux t₁) (to_aux t₂)

-- def L_abel_plus' (t₁ t₂ : preterm L_abel 0) : preterm L_abel 0 :=
-- @term_of_function L_abel 2 (abel_functions.plus : L_abel.functions 2) t₁ t₂
-- end preterm_aux

-- local infix ` +' `:100 := L_abel_plus'

-- local notation ` zero ` := (func abel_functions.zero : preterm L_abel 0)

-- section L_abel_term_biopsy

-- def sample1 : preterm L_abel 0 := (zero +' zero)

-- def sample2 : preterm L_abel 0 := zero

-- -- #reduce sample2

-- open expr
-- meta def sample2_expr : expr :=
-- mk_app (const `preterm.func list.nil) ([(const `L_abel list.nil), `(0), const `abel_functions.zero list.nil] : list expr)

-- end L_abel_term_biopsy

-- section simpler_biopsy

-- inductive my_inductive : Type
-- | a : my_inductive
-- | b : my_inductive
-- | f : my_inductive → my_inductive

-- open my_inductive
-- def sample3 : my_inductive := f a

-- open expr

-- meta def sample3_expr : expr :=
-- app (const `my_inductive.f list.nil) (const `my_inductive.a list.nil)

-- def sample3_again : my_inductive := by tactic.exact (sample3_expr)

-- example : sample3 = sample3_again := rfl

-- end simpler_biopsy

-- namespace tactic
-- namespace interactive
-- open interactive interactive.types expr

-- def my_test_term : preterm L_abel 0 := (zero +' zero)

-- end interactive
-- end tactic

-- section test
-- -- def my_term : preterm L_abel 0 := sorry

-- -- #check tactic.interactive.rcases

-- end test

-- section sample4

-- /-- Note: this is the same as `dfin` -/
-- inductive my_indexed_family : ℕ → Type u
-- | z {} : my_indexed_family 0
-- | s : ∀ {k}, my_indexed_family k → my_indexed_family (k+1)

-- -- meta example : ∀ {n}, has_reflect (my_indexed_family n)
-- -- | 0 z := `(z)

-- open my_indexed_family

-- def sample4 : my_indexed_family 1 := s z

-- -- #check tactic.eval_expr

-- end sample4

-- section sample4

-- inductive dfin'' : ℕ → Type
-- | fz {n} : dfin'' (n+1)
-- | fs {n} : dfin'' n → dfin'' (n+1)

-- inductive dfin' : ℕ → Type u
-- | gz {n} :  dfin' (n+1)
-- | gs {n} :  dfin' n → dfin' (n+1)

-- open dfin dfin'

-- meta instance dfin.reflect : ∀ {n}, has_reflect (dfin'' n)
-- | _ dfin''.fz := `(dfin''.fz)
-- | _ (dfin''.fs n) := `(dfin''.fs).subst (dfin.reflect n)

-- -- /- errors all over---why doesn't reflect like universe parameters? -/
-- -- meta instance dfin'.reflect : ∀ {n}, has_reflect (dfin' n)
-- -- | _ fz := `(fz)
-- -- | _ (fs n) := `(fs).subst (dfin'.reflect n)

-- end sample4

-- section reflect_preterm

-- /- Language with a single constant symbol -/
-- inductive L_pt_functions : ℕ → Type
-- | pt : L_pt_functions 0

-- def L_pt : Language.{0} := ⟨L_pt_functions, λ _, empty⟩

-- def pt_preterm : preterm L_pt 0 := preterm.func L_pt_functions.pt

-- meta def pt_preterm_reflected : expr :=
-- expr.mk_app (expr.const `preterm.func [level.zero]) [ (expr.const `L_pt list.nil), `(0), (expr.const `L_pt_functions.pt list.nil)]

-- set_option trace.app_builder true

-- -- meta def pt_preterm_reflected' : expr := by tactic.mk_app "preterm.func" [(expr.const `L_pt []), `(0), (expr.const `L_pt_functions.pt [])]

-- #check tactic.mk_app

-- meta def pt_preterm_reflected'' : tactic expr :=
-- tactic.to_expr ```(preterm.func L_pt_functions.pt : preterm L_pt 0)

-- def pt_preterm' : preterm L_pt 0 := by pt_preterm_reflected'' >>= tactic.exact

-- -- def pt_preterm' : preterm L_pt 0 := by tactic.exact pt_preterm_reflected

-- example : pt_preterm = pt_preterm' := rfl
--   -- infer type failed, incorrect number of universe levels

-- -- want: example : pt_preterm = pt_preterm' := rfl

-- end reflect_preterm

-- namespace hewwo
-- section reflect_preterm2
-- def L_pt.pt' : L_pt.functions 0 := L_pt_functions.pt

-- #reduce (by apply_instance : reflected L_pt.pt')
-- -- `(L_pt.pt')

-- meta def pt_preterm_reflected : tactic expr :=
-- tactic.mk_app ``preterm.func [`(L_pt.pt')]

-- def pt_preterm' : preterm L_pt 0 := by pt_preterm_reflected >>= tactic.exact

-- #eval tactic.trace (@expr.to_raw_fmt tt `(L_pt.pt'))

-- #check reflect

-- end reflect_preterm2
-- end hewwo

-- section reflect_preterm3

-- inductive L_pt_func_functions : ℕ → Type
-- | pt  : L_pt_func_functions 0
-- | foo : L_pt_func_functions 1

-- open L_pt_func_functions

-- def L_pt_func : Language.{0} :=
-- ⟨L_pt_func_functions, λ _, ulift empty⟩

-- -- def foo_pt_term : preterm L_pt_func 0 :=
-- -- preterm.app (preterm.func L_pt_func_functions.foo) (preterm.func L_pt_func_functions.pt)

-- -- def foo_pt_term_reflected : expr :=
-- -- begin
-- --   tactic.mk_app ``preterm.func [(by tactic.mk_app `preterm.func [`(L_pt_func_functions.foo)]), (by tactic.mk_app `preterm.func [`(L_pt_func_functions.pt)])]
-- -- end

-- -- def foo' : preterm L_pt_func 1 := preterm.func L_pt_func_functions.foo

-- -- #reduce (by apply_instance : reflected L_pt_func_functions.foo)

-- set_option trace.app_builder true

-- def my_foo : L_pt_func.functions 1 := L_pt_func_functions.foo

-- def my_pt : L_pt_func.functions 0 := L_pt_func_functions.pt

-- -- meta def foo_pt_term_reflected : tactic expr := tactic.mk_app ``preterm.func [`()]

-- meta def foo_pt_term_reflected' : tactic expr :=
-- do e₁ <- tactic.mk_app ``preterm.func [`(my_foo)],
--    e₂ <- tactic.mk_app ``preterm.func [`(my_pt)],
--    tactic.mk_app ``preterm.app [e₁, e₂]
-- -- #print foo_pt_term_reflected'

-- -- meta def bar : tactic expr :=
-- --   tactic.mk_app ``preterm.func [`(foo_mask)]

-- set_option trace.app_builder true

-- def foo_pt_term_reflected : preterm L_pt_func 0 := by (foo_pt_term_reflected' >>= tactic.exact)

-- -- #reduce foo_pt_term_reflected

-- end reflect_preterm3


/-
Copyright (c) 2020 Mario Carneiro. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Mario Carneiro
-/
import data.fintype.basic

/-!
# Derive handler for `fintype` instances

This file introduces a derive handler to automatically generate `fintype`
instances for structures and inductives.

## Implementation notes

To construct a fintype instance, we need 3 things:

  1. A list `l` of elements
  2. A proof that `l` has no duplicates
  3. A proof that every element in the type is in `l`

Now fintype is defined as a finset which enumerates all elements, so steps (1) and (2) are
bundled together. It is possible to use finset operations that remove duplicates to avoid the need
to prove (2), but this adds unnecessary functions to the constructed term, which makes it more
expensive to compute the list, and it also adds a dependence on decidable equality for the type,
which we want to avoid.

Because we will rely on fintype instances for constructor arguments, we can't actually build a list
directly, so (1) and (2) are necessarily somewhat intertwined. The inductive types we will be
proving instances for look something like this:

```
@[derive fintype]
inductive foo
| zero : foo
| one : bool → foo
| two : ∀ x : fin 3, bar x → foo
```

The list of elements that we generate is
```
{foo.zero}
∪ (finset.univ : bool).map (λ b, finset.one b)
∪ (finset.univ : Σ' x : fin 3, bar x).map (λ ⟨x, y⟩, finset.two x y)
```
except that instead of `∪`, that is `finset.union`, we use `finset.disj_union` which doesn't
require any deduplication, but does require a proof that the two parts of the union are disjoint.
We use `finset.cons` to append singletons like `foo.zero`.

The proofs of disjointness would be somewhat expensive since there are quadratically many of them,
so instead we use a "discriminant" function. Essentially, we define
```
def foo.enum : foo → ℕ
| foo.zero := 0
| (foo.one _) := 1
| (foo.two _ _) := 2
```
and now the existence of this function implies that foo.zero is not foo.two and so on because they
map to different natural numbers. We can prove that sets of natural numbers are mutually disjoint
more easily because they have a linear order: `0 < 1 < 2` so `0 ≠ 2`.

To package this argument up, we define `finset_above foo foo.enum n` to be a finset `s` together
with a proof that all elements `a ∈ s` have `n ≤ enum a`. Now we only have to prove that
`enum foo.zero = 0`, `enum (foo.one _) = 1`, etc. (linearly many proofs, all `rfl`) in order to
prove that all variants are mutually distinct.

We mirror the `finset.cons` and `finset.disj_union` functions into `finset_above.cons` and
`finset_above.union`, and this forms the main part of the finset construction.

This only handles distinguishing variants of a finset. Now we must enumerate the elements of a
variant, for example `{foo.one ff, foo.one tt}`, while at the same time proving that all these
elements have discriminant `1` in this case. To do that, we use the `finset_in` type, which
is a finset satisfying a property `P`, here `λ a, foo.enum a = 1`.

We could use `finset.bind` many times to construct the finset but it turns out to be somewhat
complicated to get good side goals for a naturally nodup version of `finset.bind` in the same way
as we did with `finset.cons` and `finset.union`. Instead, we tuple up all arguments into one type,
leveraging the `fintype` instance on `psigma`, and then define a map from this type to the
inductive type that untuples them and applies the constructor. The injectivity property of the
constructor ensures that this function is injective, so we can use `finset.map` to apply it. This
is the content of the constructor `finset_in.mk`.

That completes the proofs of (1) and (2). To prove (3), we perform one case analysis over the
inductive type, proving theorems like
```
foo.one a ∈ {foo.zero}
  ∪ (finset.univ : bool).map (λ b, finset.one b)
  ∪ (finset.univ : Σ' x : fin 3, bar x).map (λ ⟨x, y⟩, finset.two x y)
```
by seeking to the relevant disjunct and then supplying the constructor arguments. This part of the
proof is quadratic, but quite simple. (We could do it in `O(n log n)` if we used a balanced tree
for the unions.)

The tactics perform the following parts of this proof scheme:
* `mk_sigma` constructs the type `Γ` in `finset_in.mk`
* `mk_sigma_elim` constructs the function `f` in `finset_in.mk`
* `mk_sigma_elim_inj` proves that `f` is injective
* `mk_sigma_elim_eq` proves that `∀ a, enum (f a) = k`
* `mk_finset` constructs the finset `S = {foo.zero} ∪ ...` by recursion on the variants
* `mk_finset_total` constructs the proof `|- foo.zero ∈ S; |- foo.one a ∈ S; |- foo.two a b ∈ S`
  by recursion on the subgoals coming out of the initial `cases`
* `mk_fintype_instance` puts it all together to produce a proof of `fintype foo`.
  The construction of `foo.enum` is also done in this function.

-/

namespace derive_fintype

/-- A step in the construction of `finset.univ` for a finite inductive type.
We will set `enum` to the discriminant of the inductive type, so a `finset_above`
represents a finset that enumerates all elements in a tail of the constructor list. -/
def finset_above (α) (enum : α → ℕ) (n : ℕ) :=
{s : finset α // ∀ x ∈ s, n ≤ enum x}

/-- Construct a fintype instance from a completed `finset_above`. -/
def mk_fintype {α} (enum : α → ℕ) (s : finset_above α enum 0) (H : ∀ x, x ∈ s.1) :
  fintype α := ⟨s.1, H⟩

/-- This is the case for a simple variant (no arguments) in an inductive type. -/
def finset_above.cons {α} {enum : α → ℕ} (n)
  (a : α) (h : enum a = n) (s : finset_above α enum (n+1)) : finset_above α enum n :=
begin
  refine ⟨finset.cons a s.1 _, _⟩,
  { intro h',
    have := s.2 _ h', rw h at this,
    exact nat.not_succ_le_self n this },
  { intros x h', rcases finset.mem_cons.1 h' with rfl | h',
    { exact ge_of_eq h },
    { exact nat.le_of_succ_le (s.2 _ h') } }
end

theorem finset_above.mem_cons_self {α} {enum : α → ℕ} {n a h s} :
  a ∈ (@finset_above.cons α enum n a h s).1 := multiset.mem_cons_self _ _

theorem finset_above.mem_cons_of_mem {α} {enum : α → ℕ} {n a h s b} :
  b ∈ (s : finset_above _ _ _).1 → b ∈ (@finset_above.cons α enum n a h s).1 :=
multiset.mem_cons_of_mem

/-- The base case is when we run out of variants; we just put an empty finset at the end. -/
def finset_above.nil {α} {enum : α → ℕ} (n) : finset_above α enum n := ⟨∅, by rintro _ ⟨⟩⟩

instance (α enum n) : inhabited (finset_above α enum n) := ⟨finset_above.nil _⟩

/-- This is a finset covering a nontrivial variant (with one or more constructor arguments).
The property `P` here is `λ a, enum a = n` where `n` is the discriminant for the current
variant. -/
@[nolint has_nonempty_instance]
def finset_in {α} (P : α → Prop) := {s : finset α // ∀ x ∈ s, P x}

/-- To construct the finset, we use an injective map from the type `Γ`, which will be the
sigma over all constructor arguments. We use sigma instances and existing fintype instances
to prove that `Γ` is a fintype, and construct the function `f` that maps `⟨a, b, c, ...⟩`
to `C_n a b c ...` where `C_n` is the nth constructor, and `mem` asserts
`enum (C_n a b c ...) = n`. -/
def finset_in.mk {α} {P : α → Prop} (Γ) [fintype Γ]
  (f : Γ → α) (inj : function.injective f) (mem : ∀ x, P (f x)) : finset_in P :=
⟨finset.univ.map ⟨f, inj⟩,
 λ x h, by rcases finset.mem_map.1 h with ⟨x, _, rfl⟩; exact mem x⟩

theorem finset_in.mem_mk {α} {P : α → Prop} {Γ} {s : fintype Γ} {f : Γ → α} {inj mem a}
  (b) (H : f b = a) : a ∈ (@finset_in.mk α P Γ s f inj mem).1 :=
finset.mem_map.2 ⟨_, finset.mem_univ _, H⟩

/-- For nontrivial variants, we split the constructor list into a `finset_in` component for the
current constructor and a `finset_above` for the rest. -/
def finset_above.union {α} {enum : α → ℕ} (n)
  (s : finset_in (λ a, enum a = n)) (t : finset_above α enum (n+1)) : finset_above α enum n :=
begin
  refine ⟨finset.disj_union s.1 t.1 _, _⟩,
  { rw finset.disjoint_left,
    intros a hs ht,
    have := t.2 _ ht, rw s.2 _ hs at this,
    exact nat.not_succ_le_self n this },
  { intros x h', rcases finset.mem_disj_union.1 h' with h' | h',
    { exact ge_of_eq (s.2 _ h') },
    { exact nat.le_of_succ_le (t.2 _ h') } }
end

theorem finset_above.mem_union_left {α} {enum : α → ℕ} {n s t a}
  (H : a ∈ (s : finset_in _).1) : a ∈ (@finset_above.union α enum n s t).1 :=
multiset.mem_add.2 (or.inl H)

theorem finset_above.mem_union_right {α} {enum : α → ℕ} {n s t a}
  (H : a ∈ (t : finset_above _ _ _).1) : a ∈ (@finset_above.union α enum n s t).1 :=
multiset.mem_add.2 (or.inr H)

end derive_fintype

namespace tactic

open derive_fintype tactic expr

namespace derive_fintype

/-- Construct the term `Σ' (a:A) (b:B a) (c:C a b), unit` from
`Π (a:A) (b:B a), C a b → T` (the type of a constructor). -/
meta def mk_sigma : expr → tactic expr
| (expr.pi n bi d b) := do
  p ← mk_local' n bi d,
  e ← mk_sigma (expr.instantiate_var b p),
  tactic.mk_app ``psigma [d, bind_lambda e p]
| _ := pure `(unit)

/-- Prove the goal `(Σ' (a:A) (b:B a) (c:C a b), unit) → T`
(this is the function `f` in `finset_in.mk`) using recursive `psigma.elim`,
finishing with the constructor. The two arguments are the type of the constructor,
and the constructor term itself; as we recurse we add arguments
to the constructor application and destructure the pi type of the constructor. We return the number
of `psigma.elim` applications constructed, which is the number of constructor arguments. -/
meta def mk_sigma_elim : expr → expr → tactic ℕ
| (expr.pi n bi d b) c := do
  refine ``(@psigma.elim %%d _ _ _),
  i ← intro_fresh n,
  (+ 1) <$> mk_sigma_elim (expr.instantiate_var b i) (c i)
| _ c := do intro1, exact c $> 0

/-- Prove the goal `a, b |- f a = f b → g a = g b` where `f` is the function we constructed in
`mk_sigma_elim`, and `g` is some other term that gets built up and eventually closed by
reflexivity. Here `a` and `b` have sigma types so the proof approach is to case on `a` and `b`
until the goal reduces to `C_n a1 ... am = C_n b1 ... bm → ⟨a1, ..., am⟩ = ⟨b1, ..., bm⟩`, at which
point cases on the equality reduces the problem to reflexivity.

The arguments are the number `m` returned from `mk_sigma_elim`, and the hypotheses `a,b` that we
need to case on. -/
meta def mk_sigma_elim_inj : ℕ → expr → expr → tactic unit
| (m+1) x y := do
  [(_, [x1, x2])] ← cases x,
  [(_, [y1, y2])] ← cases y,
  mk_sigma_elim_inj m x2 y2
| 0 x y := do
  cases x, cases y,
  is ← intro1 >>= injection,
  is.mmap' cases,
  reflexivity

/-- Prove the goal `a |- enum (f a) = n`, where `f` is the function constructed in `mk_sigma_elim`,
and `enum` is a function that reduces to `n` on the constructor `C_n`. Here we just have to case on
`a` `m` times, and then `reflexivity` finishes the proof. -/
meta def mk_sigma_elim_eq : ℕ → expr → tactic unit
| (n+1) x := do
  [(_, [x1, x2])] ← cases x,
  mk_sigma_elim_eq n x2
| 0 x := reflexivity

/-- Prove the goal `|- finset_above T enum k`, where `T` is the inductive type and `enum` is the
discriminant function. The arguments are `args`, the parameters to the inductive type (and all
constructors), `k`, the index of the current variant, and `cs`, the list of constructor names.
This uses `finset_above.cons` for basic variants and `finset_above.union` for variants with
arguments, using the auxiliary functions `mk_sigma`, `mk_sigma_elim`, `mk_sigma_elim_inj`,
`mk_sigma_elim_eq` to close subgoals. -/
meta def mk_finset (ls : list level) (args : list expr) : ℕ → list name → tactic unit
| k (c::cs) := do
  let e := (expr.const c ls).mk_app args,
  t ← infer_type e,
  if is_pi t then do
    to_expr ``(finset_above.union %%(reflect k)) tt ff >>=
      (λ c, apply c {new_goals := new_goals.all}),
    Γ ← mk_sigma t,
    to_expr ``(finset_in.mk %%Γ) tt ff >>= (λ c, apply c {new_goals := new_goals.all}),
    n ← mk_sigma_elim t e,
    intro1 >>= (λ x, intro1 >>= mk_sigma_elim_inj n x),
    intro1 >>= mk_sigma_elim_eq n,
    mk_finset (k+1) cs
  else do
    c ← to_expr ``(finset_above.cons %%(reflect k) %%e) tt ff,
    apply c {new_goals := new_goals.all}, reflexivity,
    mk_finset (k+1) cs
| k [] := applyc ``finset_above.nil

/-- Prove the goal `|- Σ' (a:A) (b: B a) (c:C a b), unit` given a list of terms `a, b, c`. -/
meta def mk_sigma_mem : list expr → tactic unit
| (x::xs) := fconstructor >> exact x >> mk_sigma_mem xs
| [] := fconstructor $> ()

/-- This function is called to prove `a : T |- a ∈ S.1` where `S` is the `finset_above` constructed
by `mk_finset`, after the initial cases on `a : T`, producing a list of subgoals. For each case,
we have to navigate past all the variants that don't apply (which is what the `tac` input tactic
does), and then call either `finset_above.mem_cons_self` for trivial variants or
`finset_above.mem_union_left` and `finset_in.mem_mk` for nontrivial variants. Either way the proof
is quite simple. -/
meta def mk_finset_total : tactic unit → list (name × list expr) → tactic unit
| tac [] := done
| tac ((_, xs) :: gs) := do
  tac,
  b ← succeeds (applyc ``finset_above.mem_cons_self),
  if b then
    mk_finset_total (tac >> applyc ``finset_above.mem_cons_of_mem) gs
  else do
    applyc ``finset_above.mem_union_left,
    applyc ``finset_in.mem_mk {new_goals := new_goals.all},
    mk_sigma_mem xs,
    reflexivity,
    mk_finset_total (tac >> applyc ``finset_above.mem_union_right) gs

end derive_fintype

open tactic.derive_fintype

/-- Proves `|- fintype T` where `T` is a non-recursive inductive type with no indices,
where all arguments to all constructors are fintypes. -/
meta def mk_fintype_instance : tactic unit :=
do
  intros,
  `(fintype %%e) ← target >>= whnf,
  (const I ls, args) ← pure (get_app_fn_args e),
  env ← get_env,
  let cs := env.constructors_of I,
  guard (env.inductive_num_indices I = 0) <|>
    fail "@[derive fintype]: inductive indices are not supported",
  guard (¬ env.is_recursive I) <|>
    fail ("@[derive fintype]: recursive inductive types are " ++
          "not supported (they are also usually infinite)"),
  applyc ``mk_fintype {new_goals := new_goals.all},
  intro1 >>= cases >>= (λ gs,
    gs.enum.mmap' $ λ ⟨i, _⟩, exact (reflect i)),
  mk_finset ls args 0 cs,
  intro1 >>= cases >>= mk_finset_total skip

/--
Tries to derive a `fintype` instance for inductives and structures.

For example:
```
@[derive fintype]
inductive foo (n m : ℕ)
| zero : foo
| one : bool → foo
| two : fin n → fin m → foo
```
Here, `@[derive fintype]` adds the instance `foo.fintype`. The underlying finset
definitionally unfolds to a list that enumerates the elements of the inductive in
lexicographic order.

If the structure/inductive has a type parameter `α`, then the generated instance will have an
argument `fintype α`, even if it is not used.  (This is due to the implementation using
`instance_derive_handler`.)
-/
@[derive_handler] meta def fintype_instance : derive_handler :=
instance_derive_handler ``fintype mk_fintype_instance

end tactic


/-
Copyright (c) 2022 Andrew Yang. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Andrew Yang
-/
import category_theory.elementwise
import category_theory.adjunction.evaluation
import category_theory.sites.sheafification

/-!

# Subsheaf of types

We define the sub(pre)sheaf of a type valued presheaf.

## Main results

- `category_theory.grothendieck_topology.subpresheaf` :
  A subpresheaf of a presheaf of types.
- `category_theory.grothendieck_topology.subpresheaf.sheafify` :
  The sheafification of a subpresheaf as a subpresheaf. Note that this is a sheaf only when the
  whole sheaf is.
- `category_theory.grothendieck_topology.subpresheaf.sheafify_is_sheaf` :
  The sheafification is a sheaf
- `category_theory.grothendieck_topology.subpresheaf.sheafify_lift` :
  The descent of a map into a sheaf to the sheafification.
- `category_theory.grothendieck_topology.image_sheaf` : The image sheaf of a morphism.
- `category_theory.grothendieck_topology.image_factorization` : The image sheaf as a
  `limits.image_factorization`.
-/

universes w v u

open opposite category_theory

namespace category_theory.grothendieck_topology

variables {C : Type u} [category.{v} C] (J : grothendieck_topology C)

/-- A subpresheaf of a presheaf consists of a subset of `F.obj U` for every `U`,
compatible with the restriction maps `F.map i`. -/
@[ext]
structure subpresheaf (F : Cᵒᵖ ⥤ Type w) :=
(obj : Π U, set (F.obj U))
(map : Π {U V : Cᵒᵖ} (i : U ⟶ V), (obj U) ⊆ (F.map i) ⁻¹' (obj V))

variables {F F' F'' : Cᵒᵖ ⥤ Type w} (G G' : subpresheaf F)

instance : partial_order (subpresheaf F) :=
partial_order.lift subpresheaf.obj subpresheaf.ext

instance : has_top (subpresheaf F) :=
⟨⟨λ U, ⊤, λ U V i x h, _root_.trivial⟩⟩

instance : nonempty (subpresheaf F) := infer_instance

/-- The subpresheaf as a presheaf. -/
@[simps]
def subpresheaf.to_presheaf : Cᵒᵖ ⥤ Type w :=
{ obj := λ U, G.obj U,
  map := λ U V i x, ⟨F.map i x, G.map i x.prop⟩,
  map_id' := λ X, by { ext ⟨x, _⟩, dsimp, rw F.map_id, refl },
  map_comp' := λ X Y Z i j, by { ext ⟨x, _⟩, dsimp, rw F.map_comp, refl } }

instance {U} : has_coe (G.to_presheaf.obj U) (F.obj U) :=
coe_subtype

/-- The inclusion of a subpresheaf to the original presheaf. -/
@[simps]
def subpresheaf.ι : G.to_presheaf ⟶ F :=
{ app := λ U x, x }

instance : mono G.ι :=
⟨λ H f₁ f₂ e, nat_trans.ext f₁ f₂ $ funext $ λ U,
  funext $ λ x, subtype.ext $ congr_fun (congr_app e U) x⟩

/-- The inclusion of a subpresheaf to a larger subpresheaf -/
@[simps]
def subpresheaf.hom_of_le {G G' : subpresheaf F} (h : G ≤ G') : G.to_presheaf ⟶ G'.to_presheaf :=
{ app := λ U x, ⟨x, h U x.prop⟩ }

instance {G G' : subpresheaf F} (h : G ≤ G') : mono (subpresheaf.hom_of_le h) :=
⟨λ H f₁ f₂ e, nat_trans.ext f₁ f₂ $ funext $ λ U,
  funext $ λ x, subtype.ext $ (congr_arg subtype.val $ (congr_fun (congr_app e U) x : _) : _)⟩

@[simp, reassoc]
lemma subpresheaf.hom_of_le_ι  {G G' : subpresheaf F} (h : G ≤ G') :
  subpresheaf.hom_of_le h ≫ G'.ι = G.ι :=
by { ext, refl }

instance : is_iso (subpresheaf.ι (⊤ : subpresheaf F)) :=
begin
  apply_with nat_iso.is_iso_of_is_iso_app { instances := ff },
  { intro X, rw is_iso_iff_bijective,
    exact ⟨subtype.coe_injective, λ x, ⟨⟨x, _root_.trivial⟩, rfl⟩⟩ }
end

lemma subpresheaf.eq_top_iff_is_iso : G = ⊤ ↔ is_iso G.ι :=
begin
  split,
  { rintro rfl, apply_instance },
  { introI H, ext U x, apply (iff_true _).mpr, rw ← is_iso.inv_hom_id_apply (G.ι.app U) x,
    exact ((inv (G.ι.app U)) x).2 }
end

/-- If the image of a morphism falls in a subpresheaf, then the morphism factors through it. -/
@[simps]
def subpresheaf.lift (f : F' ⟶ F) (hf : ∀ U x, f.app U x ∈ G.obj U) : F' ⟶ G.to_presheaf :=
{ app := λ U x, ⟨f.app U x, hf U x⟩,
  naturality' := by { have := elementwise_of f.naturality, intros, ext, simp [this] } }

@[simp, reassoc]
lemma subpresheaf.lift_ι (f : F' ⟶ F) (hf : ∀ U x, f.app U x ∈ G.obj U) :
  G.lift f hf ≫ G.ι = f := by { ext, refl }

/-- Given a subpresheaf `G` of `F`, an `F`-section `s` on `U`, we may define a sieve of `U`
consisting of all `f : V ⟶ U` such that the restriction of `s` along `f` is in `G`. -/
@[simps]
def subpresheaf.sieve_of_section {U : Cᵒᵖ} (s : F.obj U) : sieve (unop U) :=
{ arrows := λ V f, F.map f.op s ∈ G.obj (op V),
  downward_closed' := λ V W i hi j,
    by { rw [op_comp, functor_to_types.map_comp_apply], exact G.map _ hi } }

/-- Given a `F`-section `s` on `U` and a subpresheaf `G`, we may define a family of elements in
`G` consisting of the restrictions of `s` -/
def subpresheaf.family_of_elements_of_section {U : Cᵒᵖ} (s : F.obj U) :
  (G.sieve_of_section s).1.family_of_elements G.to_presheaf :=
λ V i hi, ⟨F.map i.op s, hi⟩

lemma subpresheaf.family_of_elements_compatible {U : Cᵒᵖ} (s : F.obj U) :
  (G.family_of_elements_of_section s).compatible :=
begin
  intros Y₁ Y₂ Z g₁ g₂ f₁ f₂ h₁ h₂ e,
  ext1,
  change F.map g₁.op (F.map f₁.op s) = F.map g₂.op (F.map f₂.op s),
  rw [← functor_to_types.map_comp_apply, ← functor_to_types.map_comp_apply,
    ← op_comp, ← op_comp, e],
end

lemma subpresheaf.nat_trans_naturality (f : F' ⟶ G.to_presheaf) {U V : Cᵒᵖ} (i : U ⟶ V)
  (x : F'.obj U) :
  (f.app V (F'.map i x)).1 = F.map i (f.app U x).1 :=
congr_arg subtype.val (functor_to_types.naturality _ _ f i x)

include J

/-- The sheafification of a subpresheaf as a subpresheaf.
Note that this is a sheaf only when the whole presheaf is a sheaf. -/
def subpresheaf.sheafify : subpresheaf F :=
{ obj := λ U, { s | G.sieve_of_section s ∈ J (unop U) },
  map := begin
    rintros U V i s hs,
    refine J.superset_covering _ (J.pullback_stable i.unop hs),
    intros _ _ h,
    dsimp at h ⊢,
    rwa ← functor_to_types.map_comp_apply,
  end }

lemma subpresheaf.le_sheafify : G ≤ G.sheafify J :=
begin
  intros U s hs,
  change _ ∈ J _,
  convert J.top_mem _,
  rw eq_top_iff,
  rintros V i -,
  exact G.map i.op hs,
end

variable {J}

lemma subpresheaf.eq_sheafify (h : presieve.is_sheaf J F)
  (hG : presieve.is_sheaf J G.to_presheaf) : G = G.sheafify J :=
begin
  apply (G.le_sheafify J).antisymm,
  intros U s hs,
  suffices : ((hG _ hs).amalgamate _ (G.family_of_elements_compatible s)).1 = s,
  { rw ← this, exact ((hG _ hs).amalgamate _ (G.family_of_elements_compatible s)).2 },
  apply (h _ hs).is_separated_for.ext,
  intros V i hi,
  exact (congr_arg subtype.val ((hG _ hs).valid_glue (G.family_of_elements_compatible s) _ hi) : _)
end

lemma subpresheaf.sheafify_is_sheaf (hF : presieve.is_sheaf J F) :
  presieve.is_sheaf J (G.sheafify J).to_presheaf :=
begin
  intros U S hS x hx,
  let S' := sieve.bind S (λ Y f hf, G.sieve_of_section (x f hf).1),
  have := λ {V} {i : V ⟶ U} (hi : S' i), hi,
  choose W i₁ i₂ hi₂ h₁ h₂,
  dsimp [-sieve.bind_apply] at *,
  let x'' : presieve.family_of_elements F S' :=
    λ V i hi, F.map (i₁ hi).op (x _ (hi₂ hi)),
  have H : ∀ s, x.is_amalgamation s ↔ x''.is_amalgamation s.1,
  { intro s,
    split,
    { intros H V i hi,
      dsimp only [x''],
      conv_lhs { rw ← h₂ hi },
      rw ← H _ (hi₂ hi),
      exact functor_to_types.map_comp_apply F (i₂ hi).op (i₁ hi).op _ },
    { intros H V i hi,
      ext1,
      apply (hF _ (x i hi).2).is_separated_for.ext,
      intros V' i' hi',
      have hi'' : S' (i' ≫ i) := ⟨_, _, _, hi, hi', rfl⟩,
      have := H _ hi'',
      rw [op_comp, F.map_comp] at this,
      refine this.trans (congr_arg subtype.val (hx _ _ (hi₂ hi'') hi (h₂ hi''))) } },
  have : x''.compatible,
  { intros V₁ V₂ V₃ g₁ g₂ g₃ g₄ S₁ S₂ e,
    rw [← functor_to_types.map_comp_apply, ← functor_to_types.map_comp_apply],
    exact congr_arg subtype.val
      (hx (g₁ ≫ i₁ S₁) (g₂ ≫ i₁ S₂) (hi₂ S₁) (hi₂ S₂) (by simp only [category.assoc, h₂, e])) },
  obtain ⟨t, ht, ht'⟩ := hF _ (J.bind_covering hS (λ V i hi, (x i hi).2)) _ this,
  refine ⟨⟨t, _⟩, (H ⟨t, _⟩).mpr ht, λ y hy, subtype.ext (ht' _ ((H _).mp hy))⟩,
  show G.sieve_of_section t ∈ J _,
  refine J.superset_covering _ (J.bind_covering hS (λ V i hi, (x i hi).2)),
  intros V i hi,
  dsimp,
  rw ht _ hi,
  exact h₁ hi
end

lemma subpresheaf.eq_sheafify_iff (h : presieve.is_sheaf J F) :
  G = G.sheafify J ↔ presieve.is_sheaf J G.to_presheaf :=
⟨λ e, e.symm ▸ G.sheafify_is_sheaf h, G.eq_sheafify h⟩

lemma subpresheaf.is_sheaf_iff (h : presieve.is_sheaf J F) :
  presieve.is_sheaf J G.to_presheaf ↔
    ∀ U (s : F.obj U), G.sieve_of_section s ∈ J (unop U) → s ∈ G.obj U :=
begin
  rw ← G.eq_sheafify_iff h,
  change _ ↔ G.sheafify J ≤ G,
  exact ⟨eq.ge, (G.le_sheafify J).antisymm⟩
end

lemma subpresheaf.sheafify_sheafify (h : presieve.is_sheaf J F) :
  (G.sheafify J).sheafify J = G.sheafify J :=
((subpresheaf.eq_sheafify_iff _ h).mpr $ G.sheafify_is_sheaf h).symm

/-- The lift of a presheaf morphism onto the sheafification subpresheaf.  -/
noncomputable
def subpresheaf.sheafify_lift (f : G.to_presheaf ⟶ F') (h : presieve.is_sheaf J F') :
  (G.sheafify J).to_presheaf ⟶ F' :=
{ app := λ U s,
    (h _ s.prop).amalgamate _ ((G.family_of_elements_compatible ↑s).comp_presheaf_map f),
  naturality' :=
  begin
    intros U V i,
    ext s,
    apply (h _ ((subpresheaf.sheafify J G).to_presheaf.map i s).prop).is_separated_for.ext,
    intros W j hj,
    refine (presieve.is_sheaf_for.valid_glue _ _ _ hj).trans _,
    dsimp,
    conv_rhs { rw ← functor_to_types.map_comp_apply },
    change _ = F'.map (j ≫ i.unop).op _,
    refine eq.trans _ (presieve.is_sheaf_for.valid_glue _ _ _ _).symm,
    { dsimp at ⊢ hj, rwa functor_to_types.map_comp_apply },
    { dsimp [presieve.family_of_elements.comp_presheaf_map],
      congr' 1,
      ext1,
      exact (functor_to_types.map_comp_apply _ _ _ _).symm }
  end }

lemma subpresheaf.to_sheafify_lift (f : G.to_presheaf ⟶ F') (h : presieve.is_sheaf J F') :
  subpresheaf.hom_of_le (G.le_sheafify J) ≫ G.sheafify_lift f h = f :=
begin
  ext U s,
  apply (h _ ((subpresheaf.hom_of_le (G.le_sheafify J)).app U s).prop).is_separated_for.ext,
  intros V i hi,
  have := elementwise_of f.naturality,
  exact (presieve.is_sheaf_for.valid_glue _ _ _ hi).trans (this _ _)
end

lemma subpresheaf.to_sheafify_lift_unique (h : presieve.is_sheaf J F')
  (l₁ l₂ : (G.sheafify J).to_presheaf ⟶ F')
  (e : subpresheaf.hom_of_le (G.le_sheafify J) ≫ l₁ =
    subpresheaf.hom_of_le (G.le_sheafify J) ≫ l₂) : l₁ = l₂ :=
begin
  ext U ⟨s, hs⟩,
  apply (h _ hs).is_separated_for.ext,
  rintros V i hi,
  dsimp at hi,
  erw [← functor_to_types.naturality, ← functor_to_types.naturality],
  exact (congr_fun (congr_app e $ op V) ⟨_, hi⟩ : _)
end

lemma subpresheaf.sheafify_le (h : G ≤ G') (hF : presieve.is_sheaf J F)
  (hG' : presieve.is_sheaf J G'.to_presheaf) :
  G.sheafify J ≤ G' :=
begin
  intros U x hx,
  convert ((G.sheafify_lift (subpresheaf.hom_of_le h) hG').app U ⟨x, hx⟩).2,
  apply (hF _ hx).is_separated_for.ext,
  intros V i hi,
  have := congr_arg (λ f : G.to_presheaf ⟶ G'.to_presheaf, (nat_trans.app f (op V) ⟨_, hi⟩).1)
    (G.to_sheafify_lift (subpresheaf.hom_of_le h) hG'),
  convert this.symm,
  erw ← subpresheaf.nat_trans_naturality,
  refl,
end

omit J

section image

/-- The image presheaf of a morphism, whose components are the set-theoretic images. -/
@[simps]
def image_presheaf (f : F' ⟶ F) : subpresheaf F :=
{ obj := λ U, set.range (f.app U),
  map := λ U V i,
    by { rintros _ ⟨x, rfl⟩, have := elementwise_of f.naturality, exact ⟨_, this i x⟩ } }

@[simp] lemma top_subpresheaf_obj (U) : (⊤ : subpresheaf F).obj U = ⊤ := rfl

@[simp]
lemma image_presheaf_id : image_presheaf (𝟙 F) = ⊤ :=
by { ext, simp }

/-- A morphism factors through the image presheaf. -/
@[simps]
def to_image_presheaf (f : F' ⟶ F) : F' ⟶ (image_presheaf f).to_presheaf :=
(image_presheaf f).lift f (λ U x, set.mem_range_self _)

variables (J)

/-- A morphism factors through the sheafification of the image presheaf. -/
@[simps]
def to_image_presheaf_sheafify (f : F' ⟶ F) : F' ⟶ ((image_presheaf f).sheafify J).to_presheaf :=
 to_image_presheaf f ≫ subpresheaf.hom_of_le ((image_presheaf f).le_sheafify J)

variables {J}

@[simp, reassoc]
lemma to_image_presheaf_ι (f : F' ⟶ F) : to_image_presheaf f ≫ (image_presheaf f).ι = f :=
(image_presheaf f).lift_ι _ _

lemma image_presheaf_comp_le (f₁ : F ⟶ F') (f₂ : F' ⟶ F'') :
  image_presheaf (f₁ ≫ f₂) ≤ image_presheaf f₂ :=
λ U x hx, ⟨f₁.app U hx.some, hx.some_spec⟩

instance {F F' : Cᵒᵖ ⥤ Type (max v w)} (f : F ⟶ F') [hf : mono f] :
  is_iso (to_image_presheaf f) :=
begin
  apply_with nat_iso.is_iso_of_is_iso_app { instances := ff },
  intro X,
  rw is_iso_iff_bijective,
  split,
  { intros x y e,
    have := (nat_trans.mono_iff_mono_app _ _).mp hf X,
    rw mono_iff_injective at this,
    exact this (congr_arg subtype.val e : _) },
  { rintro ⟨_, ⟨x, rfl⟩⟩, exact ⟨x, rfl⟩ }
end

/-- The image sheaf of a morphism between sheaves, defined to be the sheafification of
`image_presheaf`. -/
@[simps]
def image_sheaf {F F' : Sheaf J (Type w)} (f : F ⟶ F') : Sheaf J (Type w) :=
⟨((image_presheaf f.1).sheafify J).to_presheaf,
  by { rw is_sheaf_iff_is_sheaf_of_type, apply subpresheaf.sheafify_is_sheaf,
    rw ← is_sheaf_iff_is_sheaf_of_type, exact F'.2 }⟩

/-- A morphism factors through the image sheaf. -/
@[simps]
def to_image_sheaf {F F' : Sheaf J (Type w)} (f : F ⟶ F') : F ⟶ image_sheaf f :=
⟨to_image_presheaf_sheafify J f.1⟩

/-- The inclusion of the image sheaf to the target. -/
@[simps]
def image_sheaf_ι {F F' : Sheaf J (Type w)} (f : F ⟶ F') : image_sheaf f ⟶ F' :=
⟨subpresheaf.ι _⟩

@[simp, reassoc]
lemma to_image_sheaf_ι {F F' : Sheaf J (Type w)} (f : F ⟶ F') :
  to_image_sheaf f ≫ image_sheaf_ι f = f :=
by { ext1, simp [to_image_presheaf_sheafify] }

instance {F F' : Sheaf J (Type w)} (f : F ⟶ F') : mono (image_sheaf_ι f) :=
(Sheaf_to_presheaf J _).mono_of_mono_map (by { dsimp, apply_instance })

instance {F F' : Sheaf J (Type w)} (f : F ⟶ F') : epi (to_image_sheaf f) :=
begin
  refine ⟨λ G' g₁ g₂ e, _⟩,
  ext U ⟨s, hx⟩,
  apply ((is_sheaf_iff_is_sheaf_of_type J _).mp G'.2 _ hx).is_separated_for.ext,
  rintros V i ⟨y, e'⟩,
  change (g₁.val.app _ ≫ G'.val.map _) _ = (g₂.val.app _ ≫ G'.val.map _) _,
  rw [← nat_trans.naturality, ← nat_trans.naturality],
  have E : (to_image_sheaf f).val.app (op V) y =
    (image_sheaf f).val.map i.op ⟨s, hx⟩ := subtype.ext e',
  have := congr_arg (λ f : F ⟶ G', (Sheaf.hom.val f).app _ y) e,
  dsimp at this ⊢,
  convert this; exact E.symm
end

/-- The mono factorization given by `image_sheaf` for a morphism. -/
def image_mono_factorization {F F' : Sheaf J (Type w)} (f : F ⟶ F') :
  limits.mono_factorisation f :=
{ I := image_sheaf f,
  m := image_sheaf_ι f,
  e := to_image_sheaf f }

/-- The mono factorization given by `image_sheaf` for a morphism is an image. -/
noncomputable
def image_factorization {F F' : Sheaf J (Type (max v u))} (f : F ⟶ F') :
  limits.image_factorisation f :=
{ F := image_mono_factorization f,
  is_image :=
  { lift := λ I, begin
      haveI := (Sheaf.hom.mono_iff_presheaf_mono J _ _).mp I.m_mono,
      refine ⟨subpresheaf.hom_of_le _ ≫ inv (to_image_presheaf I.m.1)⟩,
      apply subpresheaf.sheafify_le,
      { conv_lhs { rw ← I.fac }, apply image_presheaf_comp_le },
      { rw ← is_sheaf_iff_is_sheaf_of_type, exact F'.2 },
      { apply presieve.is_sheaf_iso J (as_iso $ to_image_presheaf I.m.1),
        rw ← is_sheaf_iff_is_sheaf_of_type, exact I.I.2 }
    end,
    lift_fac' := λ I, begin
      ext1,
      dsimp [image_mono_factorization],
      generalize_proofs h,
      rw [← subpresheaf.hom_of_le_ι h, category.assoc],
      congr' 1,
      rw [is_iso.inv_comp_eq, to_image_presheaf_ι],
    end } }

instance : limits.has_images (Sheaf J (Type (max v u))) :=
⟨λ _ _ f, ⟨⟨image_factorization f⟩⟩⟩

end image

end category_theory.grothendieck_topology


import Mathlib.Tactic.Linarith
import Aesop
-- A translation from SSA + Regions to Tree, with no proofs.
-- this allows for easy unfolding of the semantics.
-- If we carry around proofs of well formedness, the dependent typing
-- of the well-formedness leads to stuck terms.
-- Thus, we eschew the need to have proofs, and simply YOLO translate from
-- the origina SSA + Regions program into Tree.

namespace SSARgn2Tree

inductive VarName
| null
| x0
| x1
| x2
| x3
| x4
| x5
| x6
| x7
| y1
| y2
| y3
| y4
| z1
| z2
| z3
| z4
deriving DecidableEq

section StxSem

-- | The environment is a mapping from variable names to values.
abbrev Env (α : Type) := VarName → α

-- | Extend the environment with a new variable.
@[simp]
def Env.extend (name : VarName) (v : α) (e: Env α) : Env α :=
  fun name' => if name = name' then v else e name'

@[simp]
def Env.empty [Inhabited α] : Env α :=
  fun _ => default

@[simp]
def Env.map (f : α → β) (e : Env α) : Env β :=
  fun name => f (e name)

notation "∅" => Env.empty
notation e "[" name "↦" v "]" => Env.extend name v e

-- The input semantics given by the user.
class UserSemantics (opcode : Type)  where
  -- | Arguments given as (<args>, <rgn1>, ... <rgnN>)
  -- | Consider not allowing users to not be access all of 'Env'.
  opcodeEval (op: opcode) (vals : Int × Int) (rgns : (Int × Int → Int)) : Int

attribute [simp] UserSemantics.opcodeEval

inductive ASTKind : Type where
| O : ASTKind
| Os : ASTKind
| R : ASTKind
deriving Inhabited, DecidableEq

-- | The operations of the language.
inductive AST (opcode : Type): ASTKind → Type where
| assign (ret : VarName) (op : opcode)
  (args : VarName × VarName)
  (rgns : AST opcode .R) : AST opcode .O
| ops1 (op : AST opcode .O) :  AST opcode .Os
| opsmany (op : AST opcode .O) (ops : AST opcode .Os) : AST opcode .Os
| rgn (args : VarName × VarName) (body : AST opcode .Os) : AST opcode .R
| rgn0 : AST opcode .R

def AST.retname : AST opcode .O → VarName
| .assign (ret := ret) .. => ret

instance [Inhabited opcode] : Inhabited (AST opcode .O) where
  default := .assign .x0 default (.x0, .x0) .rgn0

instance [Inhabited opcode] : Inhabited (AST opcode .Os) where
  default := .ops1 default

instance [Inhabited opcode] : Inhabited (AST opcode .R) where
  default := .rgn (.x0, .x0) default

@[simp]
def Ops.ofList [Inhabited opcode] : List (AST opcode .O) → AST opcode .Os
| [] => panic! "need non-empty list"
| [x] => .ops1 x
| x :: xs => .opsmany x (Ops.ofList xs)

@[reducible, simp]
instance [Inhabited opcode] : Coe (List (AST opcode .O)) (AST opcode .Os) := ⟨Ops.ofList⟩

@[simp, reducible]
def ASTKind.eval : ASTKind → Type
| .O => Int
| .Os => Int
| .R => (Int × Int → Int)
-- evaluate an operation with repect to a particular user semantics.
@[simp]
def AST.eval [S: UserSemantics opcode]
  {astk: ASTKind} (e: Env Int): AST opcode astk → astk.eval × Env Int
| .assign ret op args r =>
    let (arg1, arg2) := args
    let retval := S.opcodeEval op (e arg1, e arg2) (r.eval e).fst
    (retval, e.extend ret retval)
| .ops1 op =>
   let (out, env) := op.eval e
   (out, env)
| .opsmany op ops =>
    let e' := (op.eval e).snd
    ops.eval e'
| .rgn args body =>
    (fun vals =>
      let e := Env.empty
      let e1 := e.extend args.fst vals.fst
      let e2 := e1.extend args.snd vals.snd
      let (outval, _) := body.eval e2
      outval, e)
| .rgn0 => (fun _ => default, e)

end StxSem

section Tree

inductive CtreeKind
| O -- op
| R -- region (higher order)
deriving Inhabited

-- closed trees, leaves are integers.
inductive Ctree  (opcode : Type) : CtreeKind → Type where
| binop
  (op : opcode)
  (lhs : Ctree opcode .O)
  (rhs : Ctree opcode .O)
  (rgns: Ctree opcode .R) : Ctree opcode .O
| rgn (f : Int × Int → Ctree opcode .O) : Ctree opcode .R
| rgn0 : Ctree opcode .R
| leaf (val : Int) : Ctree opcode .O

instance : Inhabited (Ctree opcode .O) where
  default := .leaf 10

instance : Inhabited (Ctree opcode .R) where
  default := .rgn0

-- convert an AST into a closed tree under the given environment
-- note that translation into a closed tree needs an environment,
-- to learn the values of variables.
-- This version has an `_` in the name since it needs a `Ctree.Env`, not an
-- `Env`. We will writea helper that converts `Env` into `Ctree.Env`.
@[simp, reducible]
def ASTKind.toCTree (opcode : Type) : ASTKind → Type
| .O => Ctree opcode .O
| .Os => Ctree opcode .O
| .R => Ctree opcode .R

@[simp]
def AST.toCtree_ {astk: ASTKind} (e : Env (Ctree opcode .O)) :
  AST opcode astk → (astk.toCTree opcode) × Env (Ctree opcode .O)
| .assign ret opcode (u, v) r =>
    let rval := (r.toCtree_ e).fst
    let t := .binop opcode (e u) (e v) rval
    (t, e.extend ret t)
| .rgn0 => (.rgn0, e)
| .rgn args body =>
    (.rgn fun vals =>
      let e := Env.empty
      let e1 := e.extend args.fst (.leaf vals.fst)
      let e2 := e1.extend args.snd (.leaf vals.snd)
      (body.toCtree_ e2).fst, e)
| .ops1 op =>
  let (val, e) :=op.toCtree_ e
  (val, e)
| .opsmany os o =>
    let (_, e) := os.toCtree_ e
    o.toCtree_ e

-- wrap every element in a (.leaf) constructor
@[simp]
def Ctree.Env.ofEnv (e: Env Int) : Env (Ctree opcode .O) :=
   fun name => .leaf (e name)

-- this converts an AST into a Ctree, given an environment
-- and an AST.
@[simp]
def Op.toCtree (a: AST opcode .O) (e: Env Int): Ctree opcode .O :=
    (a.toCtree_ (Ctree.Env.ofEnv e)).fst

@[simp]
def Ops.toCtree (a: AST opcode .Os) (e: Env Int) : Ctree opcode .O :=
    (a.toCtree_ (Ctree.Env.ofEnv e)).fst

@[simp]
def Region.toCtree (a: AST opcode .R) (e: Env Int) : Ctree opcode .R :=
    (a.toCtree_ (Ctree.Env.ofEnv e)).fst

-- evaluate a Ctree. note that this needs no environment.
@[simp]
def CtreeKind.eval : CtreeKind → Type
| .O => Int
| .R => Int × Int → Int

-- Note: is this literally "just" staging the partial evaluation against the environment?
@[simp]
def Ctree.eval [UserSemantics opcode] : Ctree opcode treek → treek.eval
| .binop o l r rs =>
  UserSemantics.opcodeEval o (l.eval, r.eval) rs.eval
| .leaf v => v
| .rgn0 => fun _ => default
| .rgn f => fun args => (f args).eval

end Tree

namespace MultipleInstructionTree
inductive Opcode
| add
| mul
| loop : Opcode
| ite : Opcode
| run : Opcode
| runnot : Opcode
| not : Opcode
| const : Int → Opcode
deriving Inhabited

def loopSemantics (n : Nat) (f : Int → Int) (v : Int) : Int :=
  match n with
  | 0 => v -- inline id
  | n + 1 => f (loopSemantics n f v) -- inline ∘

@[simp]
instance :  UserSemantics Opcode  where
  opcodeEval
  | .not, ⟨a, _⟩, _ => if a = 0 then 1 else 0
  | .mul, ⟨a, b⟩, _ => a * b
  | .add, ⟨a, b⟩, _ => a + b
  | .const i, ⟨_a, _b⟩, _ => i
  | .run, ⟨v, w⟩, r => r ⟨v, w⟩
  | .runnot, ⟨v, w⟩, r => r ⟨if v = 0 then 1 else 0, w⟩ -- execute: 'r(not v, w)'
  | .loop, ⟨n, init⟩, r => loopSemantics n.toNat (fun i => r ⟨i, 0⟩) init
  | _, _, _ => 42

def x_add_4_times_mul_val_eq (env: Env Int):
  let p : AST Opcode .Os :=
      Ops.ofList [
      .assign .x1 .add (.x0, .x0) .rgn0,
      .assign .x2 .add (.x1, .x1) .rgn0
      ]
  let q : AST Opcode .Os := Ops.ofList [
        .assign .x1 (.const 4) (.x0, .x0) .rgn0
      , .assign .x2 .mul (.x1, .x0) .rgn0
    ]
  (Ops.toCtree p env).eval = (Ops.toCtree q env).eval := by {
    simp only [Ops.ofList, AST.eval, Ops.toCtree,  AST.toCtree_];
    -- see that there are environments, which are folded away when calling
    -- Ctree.eval.
    simp[Ctree.eval];
    linarith
  }

def run_inline:
  let p : AST Opcode .R :=
    AST.rgn ⟨.x0, .x1⟩ $ Ops.ofList [
      .assign .x2 .add (.x0, .x1) .rgn0 -- x2 := x0 + x1
      ]
  let q : AST Opcode .R :=
    AST.rgn ⟨.x0, .x1⟩ $ Ops.ofList [
      .assign .x2 .run (.x0, .x1) (AST.rgn ⟨.x5, .x6⟩ (Ops.ofList [ -- x2 := run (x0, x1) { ^(x5, x6): return x5 + x6 }
        .assign .y1 .add (.x5, .x6) .rgn0
      ]))
    ]
  (Region.toCtree p Env.empty).eval = (Region.toCtree q Env.empty).eval := by {
    simp
  }

-- trying to convert an AST to a Ctree at any environment
-- is equivalent to converting it in an empty environment.
@[simp]
theorem AST.toCtree_rgn_equiv_empty (r: AST opcode .R) [Inhabited opcode] :
  (AST.toCtree_ env r).fst = (AST.toCtree_ Env.empty r).fst := by {
   simp
   cases r <;> simp;
}

def runnot_inline :
  let p : AST Opcode .R :=
    AST.rgn ⟨.x0, .x1⟩ $ Ops.ofList [
      .assign .y1 .run ⟨.x0, .x1⟩ r -- y1 := r(x0, x1)
    ]
  (Region.toCtree p Env.empty).eval = (Region.toCtree r Env.empty).eval := by {
    simp
  }

end MultipleInstructionTree

end SSARgn2Tree