import data.set

namespace probability_theory

-- TODO subtype.restrict?
def pi_subtype {α : Type*} {β : α → Type*} (mv : set α) := λ (g : Π i, β i) (i : mv), g i

@[reducible]
def pi_subtype_img {α : Type*} {β : α → Type*} (mv : set α) :=
  λ (g : set (Π i, β i)) , pi_subtype mv '' g

@[reducible]
def pi_unsubtype_img {α : Type*} {β : α → Type*} (mv : set α) :=
  λ (g : set (Π i : mv, β i)), pi_subtype mv ⁻¹' g

notation  `<[`S`]` := pi_subtype_img S
notation  `>[`S`]` := pi_unsubtype_img S

notation  `<[]` := pi_subtype_img _
notation  `>[]` := pi_unsubtype_img _

def set_to_subtype {α : Type*} (A : set α) (B : set α) : set A := λ x : A, ↑x ∈ B

def pi_set_to_subtype {α : Type*} {β : α → Type*} (A : set α) (B : set α)
  (f : Π i : B, β i) : Π i : set_to_subtype A B, β i := λ ⟨i, hi⟩, f ⟨i, hi⟩

lemma pi_set_to_subtype_def {α : Type*} {β : α → Type*} (A : set α) (B : set α)
  (f : Π i : B, β i) (i : α) (hi : i ∈ A) (hi' : i ∈ B) :
  pi_set_to_subtype A B f ⟨⟨i, hi⟩, hi'⟩ = f ⟨↑(⟨i, hi⟩ : A), hi'⟩ := rfl

lemma pi_set_to_subtype_def' {α : Type*} {β : α → Type*} (A : set α) (B : set α)
  (f : Π i : B, β i) (i : set_to_subtype A B) :
  pi_set_to_subtype A B f i = f ⟨↑(⟨i.val.val, i.val.property⟩ : A), i.property⟩ :=
begin
  obtain ⟨⟨_, _⟩, _⟩ := i,
  rw pi_set_to_subtype_def,
end

lemma pi_set_to_subtype_bijective {α : Type*} {β : α → Type*} {A : set α} {B : set α} (hAB : B ⊆ A)
  : function.bijective (@pi_set_to_subtype _ β A B) :=
begin
  constructor,
  { intros f f' hff',
    have h := congr_fun hff',
    simp_rw pi_set_to_subtype_def' at h,
    ext i,
    specialize h ⟨⟨i, hAB i.property⟩, i.property⟩,
    convert h; exact subtype.eq rfl },
  { intro f,
    refine ⟨(λ i, f ⟨⟨i, hAB i.property⟩, i.property⟩), _⟩,
    ext i,
    rw pi_set_to_subtype_def',
    congr,
    refine subtype.eq _, refine subtype.eq rfl },
end

lemma pi_set_to_subtype_img_preimage_idx {α : Type*} {β : α → Type*} {A : set α} {B : set α} (hAB : B ⊆ A) {b : B} (bs : set (β b)) :
pi_set_to_subtype A B '' ((λ (g : Π (i : B), β i), g b) ⁻¹' bs)
= (λ (g : Π (i : set_to_subtype A B), β i), g ⟨⟨b, hAB b.property⟩, b.property⟩) ⁻¹' bs :=
begin
  ext1 x,
  split,
  rintro ⟨bs', hbs', hbsx'⟩,
  change bs' _ ∈ bs at hbs',
  change _ ∈ bs,
  subst hbsx',
  change bs' ⟨_, _⟩ ∈ _,
  convert hbs',
  exact subtype.eq rfl,
  intro hx,
  refine ⟨λ b : B, x ⟨⟨b, hAB b.property⟩, b.property⟩, hx, _⟩, ext ⟨⟨_, _⟩, _⟩, refl
end

@[reducible]
def pi_unsubtype_set {α : Type*} {β : α → Type*} (A : set α) (B : set α) :
  set (Π i : B, β i) → set (Π i : A, β i)
  := λ g, >[set_to_subtype A B] (pi_set_to_subtype A B '' g)

notation `>>[`A`]` := pi_unsubtype_set A _
notation `>>[]` := pi_unsubtype_set _ _

def pi_unsubtype_set_same {α : Type*} {β : α → Type*} (A : set α) (a : set (Π i : A, β i)) :
  >>[A] a = a :=
begin
  rw pi_unsubtype_set,
  rw pi_unsubtype_img,
  change (_ ⁻¹' (_ '' a)) = a,
  -- TODO extract
  convert @set.preimage_image_eq _ _ (@pi_set_to_subtype _ β A A) a (pi_set_to_subtype_bijective rfl.subset).injective,
  ext f i,
  rw pi_set_to_subtype_def',
  change f _ = _,
  congr, refine subtype.eq rfl
end

lemma pi_subtype_ext {α : Type*} {β : α → Type*} {A : set α}
{f : Π i, β i} {g : Π i : A, β i} : pi_subtype A f = g ↔ ∀ i : A, f i = g i :=
by rw function.funext_iff; refl

--lemma pi_subtype_ext' {α : Type*} {β : α → Type*} {A : set α} {f g : Π i, β i} :
--pi_subtype A f = pi_subtype A g ↔ ∀ i ∈ A, f i = g i := sorry

lemma pi_subtype_subtype {α : Type*} {β : α → Type*} (A : set α) (B : set α)
  (x : Π i : α, β i) :
  pi_subtype (set_to_subtype A B) (pi_subtype A x) = λ (i : set_to_subtype A B), x i := rfl

lemma pi_unsubtype_union_img_def {α : Type*} {β : α → Type*} [∀ i : α, inhabited (β i)]
  (A : set α) (B : set α) (sb : set (Π i : B, β i)) : >>[A] sb = <[A] (>[B] sb) :=
begin
  simp_rw [pi_unsubtype_img, pi_subtype_img],
  refine set.subset.antisymm _ _; intros x h,
  { obtain ⟨x', h', h⟩ := h,
    classical,
    let y : Π i, β i := λ i, if h : i ∈ B then x' ⟨i, h⟩
      else if h : i ∈ A then x ⟨i, h⟩ else default,
    refine ⟨y, _, _⟩,
    change pi_subtype B y ∈ sb,
    convert h',
    all_goals {refine pi_subtype_ext.mpr _, rintro ⟨i, hi⟩},
      exact dif_pos hi,
    by_cases hi' : i ∈ B,
      convert dif_pos hi',
      exact (congr_fun h ⟨⟨_, hi⟩, hi'⟩).symm,
    exact (dif_neg hi').trans (dif_pos hi) },
  { obtain ⟨_, h', rfl⟩ := h,
    refine ⟨pi_subtype B _, h', _⟩,
    ext ⟨⟨_, _⟩, _⟩, refl }
end

def pi_subtype_subtype_subset {α : Type*} {β : α → Type*} {A : set α} {B : set α}
  (hba : B ⊆ A) (sb : set (Π i : B, β i)) :
  >[B] sb = >[A] (>[set_to_subtype A B] (pi_set_to_subtype A B '' sb)) :=
begin
  simp_rw [pi_unsubtype_img, set.preimage_preimage],
  refine set.subset.antisymm _ _; intros x h,
  { refine ⟨pi_subtype B x, h, _⟩, ext ⟨⟨_, _⟩, _⟩, refl },
  { obtain ⟨_, h', h⟩ := h,
    change pi_subtype B x ∈ sb,
    convert h',
    refine pi_subtype_ext.mpr _,
    rintro ⟨_, hi⟩,
    exact (congr_fun h ⟨⟨_, hba hi⟩, _⟩).symm }
end

lemma pi_unsubtype_union_img_inter {α : Type*} {β : α → Type*} (A : set α) (B : set α)
  (a : set (Π i : A, β i)) (b : set (Π i : B, β i)) :
  >[] (>>[A ∪ B] a ∩ >>[A ∪ B] b) = >[] a ∩ >[] b :=
begin
  simp_rw pi_unsubtype_img,
  rw set.preimage_inter,
  refine congr (congr_arg has_inter.inter _) _,
  exact (pi_subtype_subtype_subset (set.subset_union_left A B) _).symm,
  exact (pi_subtype_subtype_subset (set.subset_union_right A B) _).symm
end

end probability_theory


namespace ForIn

inductive Step.{u} (α : Type u)
| done  : α → Step α
| yield : α → Step α

class Fold.{u, v, w, z} (m : Type w → Type z) [Monad m] (α : outParam (Type u)) (s : Type v) : Type (max v u z (w+1)):=
(fold {β : Type w} (as : s) (init : β) (f : α → β → m (Step β)) : m β)

export Fold (fold)

class FoldMap.{u, w} (m : Type u → Type w) [Monad m] (s : Type u → Type u) : Type (max (u+1) w):=
(foldMap {α β : Type u} (as : s α) (init : β) (f : α → β → m (Step (α × β))) : m (s α × β))

export FoldMap (foldMap)

@[inline] instance {m} {α} [Monad m] : Fold m α (List α) :=
{ fold := fun as init f =>
    let rec @[specialize] loop
      | [], b    => pure b
      | a::as, b => do
        let s ← f a b
        (match s with
         | Step.done b     => pure b
         | Step.yield b => loop as b)
    loop as init }

@[inline] instance {m} [Monad m] : FoldMap m List :=
{ foldMap := fun as init f =>
    let rec @[specialize] loop
      | [], rs, b => pure (rs.reverse, b)
      | a::as, rs, b => do
        let s ← f a b
        (match s with
         | Step.done (a, b)     => pure ((a :: rs).reverse ++ as, b)
         | Step.yield (a, b) => loop as (a::rs) b)
    loop as [] init }

def tst1 : IO Nat :=
fold [1, 2, 3, 4, 5, 7, 8, 9, 10, 11, 12, 14] 0 fun a b =>
  if a % 2 == 0 then do
    IO.println (">> " ++ toString a ++ " " ++ toString b)
    (if b > 20 then return Step.done b
     else return Step.yield (a+b))
  else
    return Step.yield b

#eval tst1

def tst1' : IO Unit := do
let (as, b) ← foldMap [1, 2, 3, 4, 5, 7, 8, 9, 10, 11, 12, 14] 0 fun a b =>
  if a % 2 == 0 then do
    IO.println (">> " ++ toString a ++ " " ++ toString b)
    (if b > 20 then return Step.done (a, b)
     else return Step.yield (a/2, a+b))
  else
    return Step.yield (a, b)
IO.println as
IO.println b
pure ()

#eval tst1'

instance Prod.fold {m α β γ δ} [Monad m] [i₁ : Fold m α γ] [i₂ : Fold m β δ] : Fold m (α × β) (γ × δ) :=
{ fold := fun s init f =>
   Fold.fold s.1 init fun a x => do
     Fold.fold s.2 (Step.yield x) fun b x =>
        match x with
        | Step.done _ => return Step.done x
        | Step.yield x => do
          let s ← f (a, b) x
          (match s with
           | Step.done _     => return Step.done s
           | Step.yield _ => return Step.yield s) }

def tst2 (threshold : Nat) : IO Nat :=
fold ([1, 2, 3, 4, 5, 10], [10, 20, 30, 40, 50]) 0 fun (a, b) s => do
  IO.println (">> " ++ toString a ++ ", " ++ toString b ++ ", " ++ toString s)
  (if s > threshold then return Step.done s
   else return Step.yield (s+a+b))

#eval tst2 170
#eval tst2 800

structure Range :=
(lower upper : Nat)

@[inline] instance Range.fold {m} [Monad m] : Fold m Nat Range :=
{ fold := fun s init f =>
  let base := s.lower + s.upper - 2
  let rec @[specialize] loop : Nat → _ → _
    | 0,   b => pure b
    | i+1, b =>
      let j := base - i
      if j >= s.upper then return b
      else do
        let s ← f j b
        (match s with
         | Step.done b     => return b
         | Step.yield b => loop i b)
  loop (s.upper - 1) init }

@[inline] def range (a : Nat) (b : Option Nat := none) : Range :=
match b with
| none      => ⟨0, a⟩
| some b    => ⟨a, b⟩

instance : OfNat (Option Nat) :=
⟨fun n => some n⟩

def tst3 : IO Nat :=
fold (range 5 10) 0 fun i s => do
  IO.println (">> " ++ toString i)
  return Step.yield (s+i)

#eval tst3

theorem zeroLtOfLt : {a b : Nat} → a < b → 0 < b
| 0,   _, h => h
| a+1, b, h =>
  have a < b from Nat.ltTrans (Nat.ltSuccSelf _) h
  zeroLtOfLt this

@[inline] instance {m} {α} [Monad m] : Fold m α (Array α) :=
{ fold := fun as init f =>
    let rec @[specialize] loop : (i : Nat) → i ≤ as.size → _
      | 0, h, b   => pure b
      | i+1, h, b =>
        have h' : i < as.size          from Nat.ltOfLtOfLe (Nat.ltSuccSelf i) h
        have as.size - 1 < as.size     from Nat.subLt (zeroLtOfLt h') (decide! (0 < 1))
        have as.size - 1 - i < as.size from Nat.ltOfLeOfLt (Nat.subLe (as.size - 1) i) this; do
        let s ← f (as.get ⟨as.size - 1 - i, this⟩) b
        (match s with
         | Step.done b     => pure b
         | Step.yield b => loop i (Nat.leOfLt h') b)
    loop as.size (Nat.leRefl _) init }

-- set_option trace.compiler.ir.result true

def tst4 : IO Nat :=
fold (#[1, 2, 3, 4, 5] : Array Nat) 0 fun a b => do
  IO.println (">> " ++ toString a ++ " " ++ toString b)
  return Step.yield (a+b)

#eval tst4

end ForIn


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

namespace Option

/-!
# Bootstrapping theorems for Option

These are theorems used in the definitions of `Std.Data.List.Basic`.
New theorems should be added to `Std.Data.Option.Lemmas` if they are not needed by the bootstrap.
-/

@[simp] theorem getD_none : getD none a = a := rfl
@[simp] theorem getD_some : getD (some a) b = a := rfl

@[simp] theorem map_none' (f : α → β) : none.map f = none := rfl
@[simp] theorem map_some' (a) (f : α → β) : (some a).map f = some (f a) := rfl

@[simp] theorem none_bind (f : α → Option β) : none.bind f = none := rfl
@[simp] theorem some_bind (a) (f : α → Option β) : (some a).bind f = f a := rfl


/-
BinaryTools: Utilities for displaying and manipulating Binary data.
-/

/-
Simplification rules for ensuring type safety of Blake3Hash
-/
@[simp] theorem ByteArray.size_empty : ByteArray.empty.size = 0 :=
rfl

@[simp] theorem ByteArray.size_push (B : ByteArray) (a : UInt8) : (B.push a).size = B.size + 1 :=
by { cases B; simp only [ByteArray.push, ByteArray.size, Array.size_push] }

@[simp] theorem List.to_ByteArray_size : (L : List UInt8) → L.toByteArray.size = L.length
| [] => rfl
| a::l => by simp [List.toByteArray, to_ByteArray_loop_size]
where to_ByteArray_loop_size :
  (L : List UInt8) → (B : ByteArray) → (List.toByteArray.loop L B).size = L.length + B.size
| [], B => by simp [List.toByteArray.loop]
| a::l, B => by
    simp [List.toByteArray.loop, to_ByteArray_loop_size]
    rw [Nat.add_succ, Nat.succ_add]

universe u
universe v

/-
Type class for default conversion between two types.
-/
class Into (Target: Type v) (Source: Type u) :=
  (into: Source → Target)

export Into (into)

instance (A: Type u) : Into A A := ⟨id⟩

def String.toByteArray (s : String) : ByteArray :=
  (List.map
    (fun c : Char => c.toNat.toUInt8) s.toList).toByteArray

instance : Into ByteArray String := {
  into := String.toByteArray
}

namespace Alphabet
def base2: String := "01"
def base8: String := "01234567"
def base10: String := "0123456789"
def base16: String := "0123456789abcdef"
def base16upper: String := "0123456789ABCDEF"
def base32: String := "abcdefghijklmnopqrstuvwxyz234567"
def base32upper: String := "ABCDEFGHIJKLMNOPQRSTUVWXYZ234567"
def base32hex : String := "0123456789abcdefghijklmnopqrstuv"
def base32hexupper : String := "0123456789ABCDEFGHIJKLMNOPQRSTUV"
def base32z : String := "ybndrfg8ejkmcpqxot1uwisza345h769"
def base36 : String := "0123456789abcdefghijklmnopqrstuvwxyz"
def base36upper : String := "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"
def base58flickr : String :=
  "123456789abcdefghijkmnopqrstuvwxyzABCDEFGHJKLMNPQRSTUVWXYZ"
def base58btc : String :=
  "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz"
def base64 : String :=
  "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"
def base64url : String :=
  "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_"

end Alphabet

/-
Encode a ByteArray as a base64 String
-/
def toBase64 {I: Type u} [Into ByteArray I] (input : I) (pad: Bool := true) : String := do
  let input : ByteArray := Into.into input
  let x := ByteArray.size input % 3
  let mut bytes := input
  let mut str := ""
  if x == 1 then bytes := bytes.append [0x00, 0x00].toByteArray
  if x == 2 then bytes := bytes.append [0x00].toByteArray
  for i in [:(bytes.size / 3)] do
    let b0 := bytes.data[3 * i]
    let b1 := bytes.data[3 * i + 1]
    let b2 := bytes.data[3 * i + 2]
    let s0 := b0.shiftRight 2
    let s1 := UInt8.xor
      ((b0.land 0b00000011).shiftLeft 4)
      ((b1.land 0b11110000).shiftRight 4)
    let s2 := UInt8.xor
      ((b1.land 0b00001111).shiftLeft 2)
      ((b2.land 0b11000000).shiftRight 6)
    let s3 := b2.land 0b00111111
    str := str.push (Alphabet.base64.get s0.toNat)
    str := str.push (Alphabet.base64.get s1.toNat)
    str := str.push (Alphabet.base64.get s2.toNat)
    str := str.push (Alphabet.base64.get s3.toNat)
  if pad then do
    if x == 1 then
      str := str.set (str.length - 1) '='
      str := str.set (str.length - 2) '='
    if x == 2 then
      str := str.set (str.length - 1) '='
    return str
  else
    if x == 1 then str := str.dropRight 2
    if x == 2 then str := str.dropRight 1
    return str


namespace util.data.nonempty_list

universe u
variable {α : Type u}

-- structure nonempty_list (α : Type u) :=
--   (head : α)
--   (tail : list α)

inductive nonempty_list (α : Type u)
| singleton : α → nonempty_list
| cons : α → nonempty_list → nonempty_list

namespace nonempty_list

-- def singleton : α → nonempty_list α := λ x, { head := x, tail := [] }

@[reducible] def head : nonempty_list α → α
| (nonempty_list.singleton x) := x
| (nonempty_list.cons x _)    := x

@[reducible] def to_list : nonempty_list α → list α
| (singleton x) := [x]
| (cons x xs)   := x :: to_list xs

@[reducible] def tail : nonempty_list α → list α
| (singleton _) := []
| (cons _ xs)   := to_list xs

-- def append (xs : nonempty_list α) (ys : nonempty_list α) : nonempty_list α := {
--   head := xs.head,
--   tail := xs.tail ++ ys.to_list
-- }

def append : nonempty_list α → nonempty_list α → nonempty_list α
| (singleton x) ys := cons x ys
| (cons x xs)   ys := cons x (append xs ys)

instance : has_append (nonempty_list α) := ⟨ append ⟩

-- lemma append_assoc : Π (xs ys zs : nonempty_list α), append (append xs ys) zs = append xs (append ys zs) :=
-- begin
--   intros,
--   unfold append,
--   simp
-- end

@[simp] lemma singleton_append (x : α) (xs : nonempty_list α) : singleton x ++ xs = cons x xs :=
rfl

@[simp] lemma cons_append (x : α) (xs : nonempty_list α) (ys : nonempty_list α) : cons x xs ++ ys = cons x (xs ++ ys) :=
rfl

lemma append_assoc (xs ys zs : nonempty_list α) : (xs ++ ys) ++ zs = xs ++ (ys ++ zs) :=
begin
  induction xs,
  simp *,
  simp *
end

end nonempty_list

end util.data.nonempty_list


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

open Function

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

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

export LawfulFunctor (map_const id_map comp_map)

attribute [simp] id_map

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

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

export LawfulApplicative (seqLeft_eq seqRight_eq pure_seq map_pure seq_pure seq_assoc)

attribute [simp] map_pure seq_pure

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

class LawfulMonad (m : Type u → Type v) [Monad m] extends LawfulApplicative m : Prop where
  bind_pure_comp (f : α → β) (x : m α) : x >>= (fun a => pure (f a)) = f <$> x
  bind_map       {α β : Type u} (f : m (α → β)) (x : m α) : f >>= (. <$> x) = f <*> x
  pure_bind      (x : α) (f : α → m β) : pure x >>= f = f x
  bind_assoc     (x : m α) (f : α → m β) (g : β → m γ) : x >>= f >>= g = x >>= fun x => f x >>= g
  map_pure g x    := (by rw [← bind_pure_comp, pure_bind])
  seq_pure g x    := (by rw [← bind_map]; simp [map_pure, bind_pure_comp])
  seq_assoc x g h := (by simp [← bind_pure_comp, ← bind_map, bind_assoc, pure_bind])

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

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

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

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

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

@[simp] theorem bind_pure_unit [Monad m] [LawfulMonad m] {x : m PUnit} : (x >>= fun _ => pure ⟨⟩) = x := by
  rw [bind_pure]

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

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

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

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

/-! # Id -/

namespace Id

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

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

end Id

/-! # ExceptT -/

namespace ExceptT

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

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

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

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

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

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

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

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

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

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

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

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

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

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

end ExceptT

/-! # ReaderT -/

namespace ReaderT

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

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

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

@[simp] theorem run_mapConst [Monad m] (a : α) (x : ReaderT ρ m β) (ctx : ρ)
    : (Functor.mapConst a x).run ctx = Functor.mapConst a (x.run ctx) := rfl

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

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

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

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

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

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

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

instance [Monad m] [LawfulFunctor m] : LawfulFunctor (ReaderT ρ m) where
  id_map    := by intros; apply ext; simp
  map_const := by intros; funext a b; apply ext; intros; simp [map_const]
  comp_map  := by intros; apply ext; intros; simp [comp_map]

instance [Monad m] [LawfulApplicative m] : LawfulApplicative (ReaderT ρ m) where
  seqLeft_eq  := by intros; apply ext; intros; simp [seqLeft_eq]
  seqRight_eq := by intros; apply ext; intros; simp [seqRight_eq]
  pure_seq    := by intros; apply ext; intros; simp [pure_seq]
  map_pure    := by intros; apply ext; intros; simp [map_pure]
  seq_pure    := by intros; apply ext; intros; simp [seq_pure]
  seq_assoc   := by intros; apply ext; intros; simp [seq_assoc]

instance [Monad m] [LawfulMonad m] : LawfulMonad (ReaderT ρ m) where
  bind_pure_comp := by intros; apply ext; intros; simp [LawfulMonad.bind_pure_comp]
  bind_map       := by intros; apply ext; intros; simp [bind_map]
  pure_bind      := by intros; apply ext; intros; simp
  bind_assoc     := by intros; apply ext; intros; simp

end ReaderT

/-! # StateRefT -/

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

/-! # StateT -/

namespace StateT

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

end StateT


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

namespace Lean.Compiler

scoped notation:max "◾" => lcErased

namespace LCNF

def erasedExpr := mkConst ``lcErased

def _root_.Lean.Expr.isErased (e : Expr) :=
  e.isAppOf ``lcErased

def isPropFormerTypeQuick : Expr → Bool
  | .forallE _ _ b _ => isPropFormerTypeQuick b
  | .sort .zero => true
  | _ => false

/--
Return true iff `type` is `Prop` or `As → Prop`.
-/
partial def isPropFormerType (type : Expr) : MetaM Bool := do
  match isPropFormerTypeQuick type with
  | true => return true
  | false => go type #[]
where
  go (type : Expr) (xs : Array Expr) : MetaM Bool := do
    match type with
    | .sort .zero => return true
    | .forallE n d b c => Meta.withLocalDecl n c (d.instantiateRev xs) fun x => go b (xs.push x)
    | _ =>
      let type ← Meta.whnfD (type.instantiateRev xs)
      match type with
      | .sort .zero => return true
      | .forallE .. => go type #[]
      | _ => return false

/--
Return true iff `e : Prop` or `e : As → Prop`.
-/
def isPropFormer (e : Expr) : MetaM Bool := do
  isPropFormerType (← Meta.inferType e)

/-!
The code generator uses a format based on A-normal form.
This normal form uses many let-expressions and it is very convenient for
applying compiler transformations. However, it creates a few issues
in a dependently typed programming language.

- Many casts are needed.
- It is too expensive to ensure we are not losing typeability when creating join points
  and simplifying let-values
- It may not be possible to create a join point because the resulting expression is
  not type correct. For example, suppose we are trying to create a join point for
  making the following `match` terminal.
  ```
  let x := match a with | true => b | false => c;
  k[x]
  ```
  and want to transform this code into
  ```
  let jp := fun x => k[x]
  match a with
  | true => jp b
  | false => jp c
  ```
  where `jp` is a new join point (i.e., a local function that is always fully applied and
  tail recursive). In many examples in the Lean code-base, we have to skip this transformation
  because it produces a type-incorrect term. Recall that types/propositions in `k[x]` may rely on
  the fact that `x` is definitionally equal to `match a with ...` before the creation of
  the join point.

Thus, in the first code generator pass, we convert types into a `LCNFType` (Lean Compiler Normal Form Type).
The method `toLCNFType` produces a type with the following properties:

- All constants occurring in the result type are inductive datatypes.
- The arguments of type formers are type formers, or `◾`. We use `◾` to denote erased information.
- All type definitions are expanded. If reduction gets stuck, it is replaced with `◾`.

Remark: you can view `◾` occurring in a type position as the "any type".
Remark: in our runtime, `◾` is represented as `box(0)`.

The goal is to preserve as much information as possible and avoid the problems described above.
Then, we don't have `let x := v; ...` in LCNF code when `x` is a type former.
If the user provides a `let x := v; ...` where x is a type former, we can always expand it when
converting into LCNF.
Thus, given a `let x := v, ...` in occurring in LCNF, we know `x` cannot occur in any type since it is
not a type former.

We try to preserve type information because they unlock new optimizations, and we can type check
the result produced by each code generator step.

Below, we provide some example programs and their erased variants:
-- 1. Source type: `f: (n: Nat) -> (tupleN Nat n)`.
      LCNF type: `f: Nat -> ◾`.
      We convert the return type `(tupleN Nat n) to `◾`, since we cannot reduce
      `(tupleN Nat n)` to a term of the form `(InductiveTy ...)`.

-- 2. Source type: `f: (n: Nat) (fin: Fin n) -> (tupleN Nat fin)`.
      LCNF type: `f: Nat -> Fin ◾ -> ◾`.
      Since `(Fin n)` has dependency on `n`, we erase the `n` to get the
      type `(Fin ◾)`.
-/

open Meta in
/--
Convert a Lean type into a LCNF type used by the code generator.
-/
partial def toLCNFType (type : Expr) : MetaM Expr := do
  if (← isProp type) then
    return erasedExpr
  let type ← whnfEta type
  match type with
  | .sort u     => return .sort u
  | .const ..   => visitApp type #[]
  | .lam n d b bi =>
    withLocalDecl n bi d fun x => do
      let d ← toLCNFType d
      let b ← toLCNFType (b.instantiate1 x)
      if b.isErased then
        return b
      else
        return Expr.lam n d (b.abstract #[x]) bi
  | .forallE .. => visitForall type #[]
  | .app ..  => type.withApp visitApp
  | .fvar .. => visitApp type #[]
  | _        => return erasedExpr
where
  whnfEta (type : Expr) : MetaM Expr := do
    let type ← whnf type
    let type' := type.eta
    if type' != type then
      whnfEta type'
    else
      return type

  visitForall (e : Expr) (xs : Array Expr) : MetaM Expr := do
    match e with
    | .forallE n d b bi =>
      let d := d.instantiateRev xs
      withLocalDecl n bi d fun x => do
        let d := (← toLCNFType d).abstract xs
        return .forallE n d (← visitForall b (xs.push x)) bi
    | _ =>
      let e ← toLCNFType (e.instantiateRev xs)
      return e.abstract xs

  visitApp (f : Expr) (args : Array Expr) := do
    let fNew ← match f with
      | .const declName us =>
        let .inductInfo _ ← getConstInfo declName | return erasedExpr
        pure <| .const declName us
      | .fvar .. => pure f
      | _ => return erasedExpr
    let mut result := fNew
    for arg in args do
      if (← isProp arg) then
        result := mkApp result erasedExpr
      else if (← isPropFormer arg) then
        result := mkApp result erasedExpr
      else if (← isTypeFormer arg) then
        result := mkApp result (← toLCNFType arg)
      else
        result := mkApp result erasedExpr
    return result

mutual

partial def joinTypes (a b : Expr) : Expr :=
  joinTypes? a b |>.getD erasedExpr

partial def joinTypes? (a b : Expr) : Option Expr := do
  if a.isErased || b.isErased then
    return erasedExpr -- See comment at `compatibleTypes`.
  else if a == b then
    return a
  else
    let a' := a.headBeta
    let b' := b.headBeta
    if a != a' || b != b' then
      joinTypes? a' b'
    else
      match a, b with
      | .mdata _ a, b => joinTypes? a b
      | a, .mdata _ b => joinTypes? a b
      | .app f a, .app g b =>
        (do return .app (← joinTypes? f g) (← joinTypes? a b))
         <|>
        return erasedExpr
      | .forallE n d₁ b₁ _, .forallE _ d₂ b₂ _ =>
        (do return .forallE n (← joinTypes? d₁ d₂) (joinTypes b₁ b₂) .default)
        <|>
        return erasedExpr
      | .lam n d₁ b₁ _, .lam _ d₂ b₂ _ =>
        (do return .lam n (← joinTypes? d₁ d₂) (joinTypes b₁ b₂) .default)
        <|>
        return erasedExpr
      | _, _ => return erasedExpr

end

/--
Return `true` if `type` is a LCNF type former type.

Remark: This is faster than `Lean.Meta.isTypeFormer`, as this
        assumes that the input `type` is an LCNF type.
-/
partial def isTypeFormerType (type : Expr) : Bool :=
  match type.headBeta with
  | .sort .. => true
  | .forallE _ _ b _ => isTypeFormerType b
  | _ => false

/--
Given a LCNF `type` of the form `forall (a_1 : A_1) ... (a_n : A_n), B[a_1, ..., a_n]` and `p_1 : A_1, ... p_n : A_n`,
return `B[p_1, ..., p_n]`.

Remark: similar to `Meta.instantiateForall`, buf for LCNF types.
-/
def instantiateForall (type : Expr) (ps : Array Expr) : CoreM Expr :=
  go 0 type
where
  go (i : Nat) (type : Expr) : CoreM Expr :=
    if h : i < ps.size then
      if let .forallE _ _ b _ := type.headBeta then
        go (i+1) (b.instantiate1 ps[i])
      else
        throwError "invalid instantiateForall, too many parameters"
    else
      return type
termination_by go i _ => ps.size - i

/--
Return `true` if `type` is a predicate.
Examples: `Nat → Prop`, `Prop`, `Int → Bool → Prop`.
-/
partial def isPredicateType (type : Expr) : Bool :=
  match type.headBeta with
  | .sort .zero => true
  | .forallE _ _ b _ => isPredicateType b
  | _ => false

/--
Return `true` if `type` is a LCNF type former type or it is an "any" type.
This function is similar to `isTypeFormerType`, but more liberal.
For example, `isTypeFormerType` returns false for `◾` and `Nat → ◾`, but
this function returns true.
-/
partial def maybeTypeFormerType (type : Expr) : Bool :=
  match type.headBeta with
  | .sort .. => true
  | .forallE _ _ b _ => maybeTypeFormerType b
  | _ => type.isErased

/--
`isClass? type` return `some ClsName` if the LCNF `type` is an instance of the class `ClsName`.
-/
def isClass? (type : Expr) : CoreM (Option Name) := do
  let .const declName _ := type.getAppFn | return none
  if isClass (← getEnv) declName then
    return declName
  else
    return none

/--
`isArrowClass? type` return `some ClsName` if the LCNF `type` is an instance of the class `ClsName`, or
if it is arrow producing an instance of the class `ClsName`.
-/
partial def isArrowClass? (type : Expr) : CoreM (Option Name) := do
  match type.headBeta with
  | .forallE _ _ b _ => isArrowClass? b
  | _ => isClass? type

partial def getArrowArity (e : Expr) :=
  match e.headBeta with
  | .forallE _ _ b _ => getArrowArity b + 1
  | _ => 0

/-- Return `true` if `type` is an inductive datatype with 0 constructors. -/
def isInductiveWithNoCtors (type : Expr) : CoreM Bool := do
  let .const declName _ := type.getAppFn | return false
  let some (.inductInfo info) := (← getEnv).find? declName | return false
  return info.numCtors == 0

end Lean.Compiler.LCNF


import Lean

-- This is a copy of Lean's DiscrTree, but with all definitional equality checks removed.

namespace CvxLean

open Lean
open Meta

inductive Key where
  | const : Name → Nat → Key
  | fvar  : FVarId → Nat → Key
  | lit   : Literal → Key
  | star  : Key
  | other : Key
  | arrow : Key
  | proj  : Name → Nat → Key
  deriving Inhabited, BEq, Repr

protected def Key.hash : Key → UInt64
  | Key.const n a => mixHash 5237 $ mixHash (hash n) (hash a)
  | Key.fvar n a  => mixHash 3541 $ mixHash (hash n) (hash a)
  | Key.lit v     => mixHash 1879 $ hash v
  | Key.star      => 7883
  | Key.other     => 2411
  | Key.arrow     => 17
  | Key.proj s i  => mixHash 11 $ mixHash (hash s) (hash i)

instance : Hashable Key := ⟨Key.hash⟩

inductive Trie (α : Type) where
  | node (vs : Array α) (children : Array (Key × Trie α)) : Trie α

structure DiscrTree (α : Type) where
  root : Std.PersistentHashMap Key (Trie α) := {}

def Key.ctorIdx : Key → Nat
  | Key.star     => 0
  | Key.other    => 1
  | Key.lit ..   => 2
  | Key.fvar ..  => 3
  | Key.const .. => 4
  | Key.arrow    => 5
  | Key.proj ..  => 6

def Key.lt : Key → Key → Bool
  | Key.lit v₁,      Key.lit v₂      => v₁ < v₂
  | Key.fvar n₁ a₁,  Key.fvar n₂ a₂  => Name.quickLt n₁.name n₂.name || (n₁ == n₂ && a₁ < a₂)
  | Key.const n₁ a₁, Key.const n₂ a₂ => Name.quickLt n₁ n₂ || (n₁ == n₂ && a₁ < a₂)
  | Key.proj s₁ i₁,  Key.proj s₂ i₂  => Name.quickLt s₁ s₂ || (s₁ == s₂ && i₁ < i₂)
  | k₁,              k₂              => k₁.ctorIdx < k₂.ctorIdx

instance : LT Key := ⟨fun a b => Key.lt a b⟩
instance (a b : Key) : Decidable (a < b) := inferInstanceAs (Decidable (Key.lt a b))

def Key.format : Key → Format
  | Key.star                   => "*"
  | Key.other                  => "◾"
  | Key.lit (Literal.natVal v) => Std.format v
  | Key.lit (Literal.strVal v) => repr v
  | Key.const k _              => Std.format k
  | Key.proj s i               => Std.format s ++ "." ++ Std.format i
  | Key.fvar k _               => Std.format k.name
  | Key.arrow                  => "→"

instance : ToFormat Key := ⟨Key.format⟩

def Key.arity : Key → Nat
  | Key.const _ a => a
  | Key.fvar _ a  => a
  | Key.arrow     => 2
  | Key.proj ..   => 1
  | _             => 0

instance : Inhabited (Trie α) := ⟨Trie.node #[] #[]⟩

namespace DiscrTree

def empty : DiscrTree α := { root := {} }

partial def Trie.format [ToMessageData α] : Trie α → MessageData
  | Trie.node vs cs => MessageData.group $ MessageData.paren $
    "node" ++ (if vs.isEmpty then MessageData.nil else " " ++ toMessageData vs)
    ++ MessageData.joinSep (cs.toList.map $ fun ⟨k, c⟩ => MessageData.paren (toMessageData k ++ " => " ++ format c)) ","

instance [ToMessageData α] : ToMessageData (Trie α) := ⟨Trie.format⟩

partial def format [ToMessageData α] (d : DiscrTree α) : MessageData :=
  let (_, r) := d.root.foldl
    (fun (p : Bool × MessageData) k c =>
      (false, p.2 ++ MessageData.paren (toMessageData k ++ " => " ++ toMessageData c)))
    (true, Format.nil)
  MessageData.group r

instance [ToMessageData α] : ToMessageData (DiscrTree α) := ⟨fun dt => format dt⟩

/- The discrimination tree ignores some implicit arguments and proofs.
   We use the following auxiliary id as a "mark". -/
private def tmpMVarId : MVarId := { name := `_discr_tree_tmp }
private def tmpStar := mkMVar tmpMVarId

instance : Inhabited (DiscrTree α) where
  default := {}

/--
  Return true iff the argument should be treated as a "wildcard" by the discrimination tree.
  - We ignore proofs because of proof irrelevance. It doesn't make sense to try to
    index their structure.
  - We ignore instance implicit arguments (e.g., `[Add α]`) because they are "morally" canonical.
    Moreover, we may have many definitionally equal terms floating around.
    Example: `Ring.hasAdd Int Int.isRing` and `Int.hasAdd`.
  - We considered ignoring implicit arguments (e.g., `{α : Type}`) since users don't "see" them,
    and may not even understand why some simplification rule is not firing.
    However, in type class resolution, we have instance such as `Decidable (@Eq Nat x y)`,
    where `Nat` is an implicit argument. Thus, we would add the path
    ```
    Decidable -> Eq -> * -> * -> * -> [Nat.decEq]
    ```
    to the discrimination tree IF we ignored the implict `Nat` argument.
    This would be BAD since **ALL** decidable equality instances would be in the same path.
    So, we index implicit arguments if they are types.
    This setting seems sensible for simplification lemmas such as:
    ```
    forall (x y : Unit), (@Eq Unit x y) = true
    ```
    If we ignore the implicit argument `Unit`, the `DiscrTree` will say it is a candidate
    simplification lemma for any equality in our goal.
  Remark: if users have problems with the solution above, we may provide a `noIndexing` annotation,
  and `ignoreArg` would return true for any term of the form `noIndexing t`.
-/
private def ignoreArg (a : Expr) (i : Nat) (infos : Array Meta.ParamInfo) : MetaM Bool := do
  if h : i < infos.size then
    let info := infos.get ⟨i, h⟩
    if info.isInstImplicit then
      return true
    else if info.isImplicit || info.isStrictImplicit then
      return not (← isType a)
    else
      isProof a
  else
    isProof a

private partial def pushArgsAux (infos : Array Meta.ParamInfo) : Nat → Expr → Array Expr → MetaM (Array Expr)
  | i, Expr.app f a, todo => do
    if (← ignoreArg a i infos) then
      pushArgsAux infos (i-1) f (todo.push tmpStar)
    else
      pushArgsAux infos (i-1) f (todo.push a)
  | _, _, todo => return todo

def mkNoindexAnnotation (e : Expr) : Expr :=
  mkAnnotation `noindex e

def hasNoindexAnnotation (e : Expr) : Bool :=
  annotation? `noindex e |>.isSome

private def pushArgs (root : Bool) (todo : Array Expr) (e : Expr) : MetaM (Key × Array Expr) := do
  if hasNoindexAnnotation e then
    return (Key.star, todo)
  else
    let fn := e.getAppFn
    let push (k : Key) (nargs : Nat) : MetaM (Key × Array Expr) := do
      let info ← getFunInfoNArgs fn nargs
      let todo ← pushArgsAux info.paramInfo (nargs-1) e todo
      return (k, todo)
    match fn with
    | Expr.lit v         => return (Key.lit v, todo)
    | Expr.const c _     =>
      let nargs := e.getAppNumArgs
      push (Key.const c nargs) nargs
    | Expr.proj s i a .. =>
      return (Key.proj s i, todo.push a)
    | Expr.fvar fvarId   =>
      let nargs := e.getAppNumArgs
      push (Key.fvar fvarId nargs) nargs
    | Expr.mvar mvarId   =>
      if mvarId == tmpMVarId then
        -- We use `tmp to mark some implicit arguments and proofs
        return (Key.star, todo)
      else
        return (Key.star, todo)
    | Expr.forallE _ d b _ =>
      if b.hasLooseBVars then
        return (Key.other, todo)
      else
        return (Key.arrow, todo.push d |>.push b)
    | _ =>
      return (Key.other, todo)

partial def mkPathAux (root : Bool) (todo : Array Expr) (keys : Array Key) : MetaM (Array Key) := do
  if todo.isEmpty then
    return keys
  else
    let e    := todo.back
    let todo := todo.pop
    let (k, todo) ← pushArgs root todo e
    mkPathAux false todo (keys.push k)

private def initCapacity := 8

def mkPath (e : Expr) : MetaM (Array Key) := do
  let todo : Array Expr := Array.mkEmpty initCapacity
  let keys : Array Key  := Array.mkEmpty initCapacity
  mkPathAux (root := true) (todo.push e) keys

private partial def createNodes (keys : Array Key) (v : α) (i : Nat) : Trie α :=
  if h : i < keys.size then
    let k := keys.get ⟨i, h⟩
    let c := createNodes keys v (i+1)
    Trie.node #[] #[(k, c)]
  else
    Trie.node #[v] #[]

private def insertVal [BEq α] (vs : Array α) (v : α) : Array α :=
  if vs.contains v then vs else vs.push v

private partial def insertAux [BEq α] (keys : Array Key) (v : α) : Nat → Trie α → Trie α
  | i, Trie.node vs cs =>
    if h : i < keys.size then
      let k := keys.get ⟨i, h⟩
      let c := Id.run $ cs.binInsertM
          (fun a b => a.1 < b.1)
          (fun ⟨_, s⟩ => let c := insertAux keys v (i+1) s; (k, c)) -- merge with existing
          (fun _ => let c := createNodes keys v (i+1); (k, c))
          (k, default)
      Trie.node vs c
    else
      Trie.node (insertVal vs v) cs

def insertCore [BEq α] (d : DiscrTree α) (keys : Array Key) (v : α) : DiscrTree α :=
  if keys.isEmpty then panic! "invalid key sequence"
  else
    let k := keys[0]!
    match d.root.find? k with
    | none =>
      let c := createNodes keys v 1
      { root := d.root.insert k c }
    | some c =>
      let c := insertAux keys v 1 c
      { root := d.root.insert k c }

def insert [BEq α] (d : DiscrTree α) (e : Expr) (v : α) : MetaM (DiscrTree α) := do
  let keys ← mkPath e
  return d.insertCore keys v

private def getKeyArgs (e : Expr) (isMatch root : Bool) : MetaM (Key × Array Expr) := do
  match e.getAppFn with
  | Expr.lit v         => return (Key.lit v, #[])
  | Expr.const c _     =>
    let nargs := e.getAppNumArgs
    return (Key.const c nargs, e.getAppRevArgs)
  | Expr.fvar fvarId   =>
    let nargs := e.getAppNumArgs
    return (Key.fvar fvarId nargs, e.getAppRevArgs)
  | Expr.mvar mvarId   =>
    if isMatch then
      return (Key.other, #[])
    else do
      return (Key.star, #[])
  | Expr.proj s i a .. =>
    return (Key.proj s i, #[a])
  | Expr.forallE _ d b _ =>
    if b.hasLooseBVars then
      return (Key.other, #[])
    else
      return (Key.arrow, #[d, b])
  | _ =>
    return (Key.other, #[])

private abbrev getMatchKeyArgs (e : Expr) (root : Bool) : MetaM (Key × Array Expr) :=
  getKeyArgs e (isMatch := true) (root := root)

private abbrev getUnifyKeyArgs (e : Expr) (root : Bool) : MetaM (Key × Array Expr) :=
  getKeyArgs e (isMatch := false) (root := root)

private def getStarResult (d : DiscrTree α) : Array α :=
  let result : Array α := Array.mkEmpty initCapacity
  match d.root.find? Key.star with
  | none                  => result
  | some (Trie.node vs _) => result ++ vs

private abbrev findKey (cs : Array (Key × Trie α)) (k : Key) : Option (Key × Trie α) :=
  cs.binSearch (k, default) (fun a b => a.1 < b.1)

private partial def getMatchLoop (todo : Array Expr) (c : Trie α) (result : Array α) : MetaM (Array α) := do
  match c with
  | Trie.node vs cs =>
    if todo.isEmpty then
      return result ++ vs
    else if cs.isEmpty then
      return result
    else
      let e     := todo.back
      let todo  := todo.pop
      let first := cs[0]! /- Recall that `Key.star` is the minimal key -/
      let (k, args) ← getMatchKeyArgs e (root := false)
      /- We must always visit `Key.star` edges since they are wildcards.
         Thus, `todo` is not used linearly when there is `Key.star` edge
         and there is an edge for `k` and `k != Key.star`. -/
      let visitStar (result : Array α) : MetaM (Array α) :=
        if first.1 == Key.star then
          getMatchLoop todo first.2 result
        else
          return result
      let visitNonStar (k : Key) (args : Array Expr) (result : Array α) : MetaM (Array α) :=
        match findKey cs k with
        | none   => pure result
        | some c => getMatchLoop (todo ++ args) c.2 result
      let result ← visitStar result
      match k with
      | Key.star  => pure result
      /-
        Recall that dependent arrows are `(Key.other, #[])`, and non-dependent arrows are `(Key.arrow, #[a, b])`.
        A non-dependent arrow may be an instance of a dependent arrow (stored at `DiscrTree`). Thus, we also visit the `Key.other` child.
      -/
      | Key.arrow => visitNonStar Key.other #[] (← visitNonStar k args result)
      | _         => visitNonStar k args result

private def getMatchRoot (d : DiscrTree α) (k : Key) (args : Array Expr) (result : Array α) : MetaM (Array α) :=
  match d.root.find? k with
  | none   => return result
  | some c => getMatchLoop args c result

/--
  Find values that match `e` in `d`.
-/
partial def getMatch (d : DiscrTree α) (e : Expr) : MetaM (Array α) := do
  Core.checkMaxHeartbeats "getMatch"
  let result := getStarResult d
  let (k, args) ← getMatchKeyArgs e (root := true)
  match k with
  | Key.star => return result
  | _        => getMatchRoot d k args result

partial def getUnify (d : DiscrTree α) (e : Expr) : MetaM (Array α) := do
  Core.checkMaxHeartbeats "getUnify"
  let (k, args) ← getUnifyKeyArgs e (root := true)
  match k with
  | Key.star => d.root.foldlM (init := #[]) fun result k c => process k.arity #[] c result
  | _ =>
    let result := getStarResult d
    match d.root.find? k with
    | none   => return result
    | some c => process 0 args c result
where
  process (skip : Nat) (todo : Array Expr) (c : Trie α) (result : Array α) : MetaM (Array α) := do
    match skip, c with
    | skip+1, Trie.node vs cs =>
      if cs.isEmpty then
        return result
      else
        cs.foldlM (init := result) fun result ⟨k, c⟩ => process (skip + k.arity) todo c result
    | 0, Trie.node vs cs => do
      if todo.isEmpty then
        return result ++ vs
      else if cs.isEmpty then
        return result
      else
        let e     := todo.back
        let todo  := todo.pop
        let (k, args) ← getUnifyKeyArgs e (root := false)
        let visitStar (result : Array α) : MetaM (Array α) :=
          let first := cs[0]!
          if first.1 == Key.star then
            process 0 todo first.2 result
          else
            return result
        let visitNonStar (k : Key) (args : Array Expr) (result : Array α) : MetaM (Array α) :=
          match findKey cs k with
          | none   => pure result
          | some c => process 0 (todo ++ args) c.2 result
        match k with
        | Key.star  => cs.foldlM (init := result) fun result ⟨k, c⟩ => process k.arity todo c result
        -- See comment a `getMatch` regarding non-dependent arrows vs dependent arrows
        | Key.arrow => visitNonStar Key.other #[] (← visitNonStar k args (← visitStar result))
        | _         => visitNonStar k args (← visitStar result)

end DiscrTree

-- From AESOP:
namespace Trie

unsafe def foldMUnsafe [Monad m] (initialKeys : Array Key)
    (f : σ → Array Key → α → m σ) (init : σ) : Trie α → m σ
  | Trie.node vs children => do
    let s ← vs.foldlM (init := init) λ s v => f s initialKeys v
    children.foldlM (init := s) λ s (k, t) =>
      t.foldMUnsafe (initialKeys.push k) f s

@[implementedBy foldMUnsafe]
opaque foldM [Monad m] (initalKeys : Array Key)
    (f : σ → Array Key → α → m σ) (init : σ) (t : Trie α) : m σ :=
  pure init

@[inline]
def fold (initialKeys : Array Key) (f : σ → Array Key → α → σ) (init : σ)
    (t : Trie α) : σ :=
  Id.run $ t.foldM initialKeys (init := init) λ s k a => return f s k a

end Trie

namespace DiscrTree

@[inline]
def foldM [Monad m] (f : σ → Array Key → α → m σ) (init : σ) (t : DiscrTree α) :
    m σ :=
  t.root.foldlM (init := init) λ s k t => t.foldM #[k] (init := s) f

@[inline]
def fold (f : σ → Array Key → α → σ) (init : σ) (t : DiscrTree α) : σ :=
  Id.run $ t.foldM (init := init) λ s keys a => return f s keys a

-- TODO inefficient since it doesn't take advantage of the Trie structure at all
@[inline]
def merge [BEq α] (t u : DiscrTree α) : DiscrTree α :=
  if t.root.size < u.root.size then loop t u else loop u t
  where
    @[inline]
    loop t u := t.fold (init := u) DiscrTree.insertCore

def values (t : DiscrTree α) : Array α :=
  t.fold (init := #[]) λ as _ a => as.push a

def toArray (t : DiscrTree α) : Array (Array Key × α) :=
  t.fold (init := #[]) λ as keys a => as.push (keys, a)

end DiscrTree

end CvxLean


import Std.Tactic.Basic
import Std.Tactic.GuardExpr
import Std.Tactic.RCases

set_option linter.missingDocs false

example (x : α × β × γ) : True := by
  rcases x with ⟨a, b, c⟩
  guard_hyp a : α
  guard_hyp b : β
  guard_hyp c : γ
  trivial

example (x : α × β × γ) : True := by
  rcases x with ⟨(a : α) : id α, -, c : id γ⟩
  guard_hyp a : α
  fail_if_success have : β := by assumption
  guard_hyp c : id γ
  trivial

example (x : (α × β) × γ) : True := by
  fail_if_success rcases x with ⟨_a, b, c⟩
  fail_if_success rcases x with ⟨⟨a:β, b⟩, c⟩
  rcases x with ⟨⟨a:α, b⟩, c⟩
  guard_hyp a : α
  guard_hyp b : β
  guard_hyp c : γ
  trivial

example : @Inhabited.{1} α × Option β ⊕ γ → True := by
  rintro (⟨⟨a⟩, _ | b⟩ | c)
  · guard_hyp a : α; trivial
  · guard_hyp a : α; guard_hyp b : β; trivial
  · guard_hyp c : γ; trivial

example : cond false Nat Int → cond true Int Nat → Nat ⊕ Unit → True := by
  rintro (x y : Int) (z | u)
  · guard_hyp x : Int; guard_hyp y : Int; guard_hyp z : Nat; trivial
  · guard_hyp x : Int; guard_hyp y : Int; guard_hyp u : Unit; trivial

example (x y : Nat) (h : x = y) : True := by
  rcases x with _|⟨⟩|z
  · guard_hyp h : Nat.zero = y; trivial
  · guard_hyp h : Nat.succ Nat.zero = y; trivial
  · guard_hyp z : Nat
    guard_hyp h : Nat.succ (Nat.succ z) = y; trivial

example (h : x = 3) (h₂ : x < 4) : x < 4 := by
  rcases h with ⟨⟩
  guard_hyp h₂ : 3 < 4; guard_target = 3 < 4; exact h₂

example (h : x = 3) (h₂ : x < 4) : x < 4 := by
  rcases h with rfl
  guard_hyp h₂ : 3 < 4; guard_target = 3 < 4; exact h₂

example (h : 3 = x) (h₂ : x < 4) : x < 4 := by
  rcases h with ⟨⟩
  guard_hyp h₂ : 3 < 4; guard_target = 3 < 4; exact h₂

example (h : 3 = x) (h₂ : x < 4) : x < 4 := by
  rcases h with rfl
  guard_hyp h₂ : 3 < 4; guard_target = 3 < 4; exact h₂

example (s : α ⊕ Empty) : True := by
  rcases s with s|⟨⟨⟩⟩
  guard_hyp s : α; trivial

example : True := by
  obtain ⟨n : Nat, _h : n = n, -⟩ : ∃ n : Nat, n = n ∧ True
  · exact ⟨0, rfl, trivial⟩
  trivial

example : True := by
  obtain (h : True) | ⟨⟨⟩⟩ : True ∨ False
  · exact Or.inl trivial
  guard_hyp h : True; trivial

example : True := by
  obtain h | ⟨⟨⟩⟩ : True ∨ False := Or.inl trivial
  guard_hyp h : True; trivial

example : True := by
  obtain ⟨h, h2⟩ := And.intro trivial trivial
  guard_hyp h : True; guard_hyp h2 : True; trivial

example : True := by
  fail_if_success obtain ⟨h, h2⟩
  trivial

example (x y : α × β) : True := by
  rcases x, y with ⟨⟨a, b⟩, c, d⟩
  guard_hyp a : α; guard_hyp b : β
  guard_hyp c : α; guard_hyp d : β
  trivial

example (x y : α ⊕ β) : True := by
  rcases x, y with ⟨a|b, c|d⟩
  · guard_hyp a : α; guard_hyp c : α; trivial
  · guard_hyp a : α; guard_hyp d : β; trivial
  · guard_hyp b : β; guard_hyp c : α; trivial
  · guard_hyp b : β; guard_hyp d : β; trivial

example (i j : Nat) : (Σ' x, i ≤ x ∧ x ≤ j) → i ≤ j := by
  intro h
  rcases h' : h with ⟨x, h₀, h₁⟩
  guard_hyp h' : h = ⟨x, h₀, h₁⟩
  apply Nat.le_trans h₀ h₁

example (x : Quot fun _ _ : α => True) (h : x = x): x = x := by
  rcases x with ⟨z⟩
  guard_hyp z : α
  guard_hyp h : Quot.mk (fun _ _ => True) z = Quot.mk (fun _ _ => True) z
  guard_target = Quot.mk (fun _ _ => True) z = Quot.mk (fun _ _ => True) z
  exact h

example (n : Nat) : True := by
  obtain _one_lt_n | _n_le_one : 1 < n + 1 ∨ n + 1 ≤ 1 := Nat.lt_or_ge 1 (n + 1)
  {trivial}; trivial

example (n : Nat) : True := by
  obtain _one_lt_n | (_n_le_one : n + 1 ≤ 1) := Nat.lt_or_ge 1 (n + 1)
  {trivial}; trivial

open Lean Elab Tactic in
/-- Asserts that the goal has `n` hypotheses. Used for testing. -/
elab "check_num_hyps " n:num : tactic => liftMetaMAtMain fun _ => do
  -- +1 because the _example recursion decl is in the list
  guard $ (← getLCtx).foldl (fun i _ => i+1) 0 = n.1.toNat + 1

example (h : ∃ x : Nat, x = x ∧ 1 = 1) : True := by
  rcases h with ⟨-, _⟩
  check_num_hyps 0
  trivial

example (h : ∃ x : Nat, x = x ∧ 1 = 1) : True := by
  rcases h with ⟨-, _, h⟩
  check_num_hyps 1
  guard_hyp h : 1 = 1
  trivial

example (h : True ∨ True ∨ True) : True := by
  rcases h with - | - | -
  iterate 3 · check_num_hyps 0; trivial

example : Bool → False → True
| false => by rintro ⟨⟩
| true => by rintro ⟨⟩

example : (b : Bool) → cond b False False → True := by
  rintro ⟨⟩ ⟨⟩

structure Baz {α : Type _} (f : α → α) : Prop where
  [inst : Nonempty α]
  h : f ∘ f = id

example {α} (f : α → α) (h : Baz f) : True := by rcases h with ⟨_⟩; trivial

example {α} (f : α → α) (h : Baz f) : True := by rcases h with @⟨_, _⟩; trivial

inductive Test : Nat → Prop
  | a (n) : Test (2 + n)
  | b {n} : n > 5 → Test (n * n)

example {n} (h : Test n) : n = n := by
  have : True := by
    rcases h with (a | b)
    · guard_hyp a : Nat
      trivial
    · guard_hyp b : ‹Nat› > 5
      trivial
  · rcases h with (a | @⟨n, b⟩)
    · guard_hyp a : Nat
      trivial
    · guard_hyp b : n > 5
      trivial

example (h : a ≤ 2 ∨ 2 < a) : True := by
  obtain ha1 | ha2 : a ≤ 2 ∨ 3 ≤ a := h
  · guard_hyp ha1 : a ≤ 2; trivial
  · guard_hyp ha2 : 3 ≤ a; trivial

example (h : a ≤ 2 ∨ 2 < a) : True := by
  obtain ha1 | ha2 : a ≤ 2 ∨ 3 ≤ a := id h
  · guard_hyp ha1 : a ≤ 2; trivial
  · guard_hyp ha2 : 3 ≤ a; trivial

inductive BaseType : Type where
  | one

inductive BaseTypeHom : BaseType → BaseType → Type where
  | loop : BaseTypeHom one one
  | id (X : BaseType) : BaseTypeHom X X

example : BaseTypeHom one one → Unit := by rintro ⟨_⟩ <;> constructor


import tactic.default

universes u v
-- #check @option.elim
-- def option.elim {α β} (x : β) (f : α → β) : option α → β
-- | none := x
-- | (some y) := f y
-- #exit
namespace list

def mfilter_map {m} [applicative m] {α β} (f : α → m (option β)) : list α → m (list.{v} β)
| [] := pure []
| (x :: xs) :=
  (λ a, option.elim a id (::)) <$> f x <*> mfilter_map xs

def mfilter_map' {m} [applicative m] [alternative m] {α β} (f : α → m β) : list α → m (list.{v} β)
| [] := pure []
| (x :: xs) :=
  ((::) <$> f x <|> pure id) <*> mfilter_map' xs

end list

namespace tactic

open native

meta def rename_many' (renames : name_map name) (strict := tt) (use_unique_names := ff)
: tactic (list (name × expr)) :=
do let hyp_name : expr → name :=
     if use_unique_names then expr.local_uniq_name else expr.local_pp_name,
   ctx ← revertible_local_context,
   -- The part of the context after (but including) the first hypthesis that
   -- must be renamed.
   let ctx_suffix := ctx.drop_while (λ h, (renames.find $ hyp_name h).is_none),
   when strict $ do {
     let ctx_names := rb_map.set_of_list (ctx_suffix.map hyp_name),
     let invalid_renames :=
       (renames.to_list.map prod.fst).filter (λ h, ¬ ctx_names.contains h),
     when ¬ invalid_renames.empty $ fail $ format.join
       [ "Cannot rename these hypotheses:\n"
       , format.join $ (invalid_renames.map to_fmt).intersperse ", "
       , format.line
       , "This is because these hypotheses either do not occur in the\n"
       , "context or they occur before a frozen local instance.\n"
       , "In the latter case, try `tactic.unfreeze_local_instances`."
       ]
   },
   -- The new names for all hypotheses in ctx_suffix.
   let new_names :=
     ctx_suffix.map $ λ h,
       (renames.find $ hyp_name h).get_or_else h.local_pp_name,
   revert_lst ctx_suffix,
   -- trace_state,
   xs ← intro_lst new_names,
   xs' ← xs.mmap infer_type,
   -- trace $ ctx_suffix.zip $ xs.zip xs',
   pure $ (ctx_suffix.map expr.local_uniq_name).zip xs

meta def find_hyp (pat : pexpr) (f : expr → tactic unit) : tactic unit := do
ls ← local_context,
pat ← pexpr_to_pattern pat,
ls.mfirst $ λ h, do
  t ← infer_type h,
  match_pattern pat t,
  f h

meta def find_all_hyps (pat : pexpr) (f : expr → tactic unit) : tactic unit := do
ls ← local_context,
pat ← pexpr_to_pattern pat,
ls.mmap' $ λ h, try $ do
  t ← infer_type h,
  match_pattern pat t,
  f h

namespace interactive
setup_tactic_parser

meta def find_hyp (id : parse ident) (_ : parse (tk ":=")) (pat : parse texpr)
  (_ : parse (tk "then")) (tac : itactic) : tactic unit := do
ls ← local_context,
pat ← pexpr_to_pattern pat,
ls.mfirst $ λ h, do
  t ← infer_type h,
  match_pattern pat t,
  rename_many (native.rb_map.of_list [(h.local_uniq_name, id)]) tt tt,
  tac,
  rename_many $ native.rb_map.of_list [(id, h.local_pp_name)]
-- #exit

meta def find_all_hyps (id : parse ident) (_ : parse (tk ":=")) (pat : parse texpr)
  (_ : parse (tk "then")) (tac : itactic) : tactic unit := do
ls ← local_context,
pat ← pexpr_to_pattern pat,
ls.reverse.mmap' (λ h,
  -- trace σ,
  -- let h := h.instantiate_locals $ list.reverse σ,
  try_or_report_error $ do {
    t ← infer_type h,
    match_pattern pat t,
    -- trace!"{h} : {t}",
    n ← tactic.revert h,
    intro id,
    intron $ n-1,
    -- xs ← rename_many' (native.rb_map.of_list [(h.local_uniq_name, id)]) tt tt,
    tac,
    h' ← get_local id,
    n ← tactic.revert h',
    intro h.local_pp_name,
    intron $ n-1,
    skip }),
skip

meta def match_le_or_lt : expr → option (expr × expr)
| `(%%x < %%y) := pure (x, y)
| `(%%x ≤ %%y) := pure (x, y)
| `(%%x > %%y) := pure (y, x)
| `(%%x ≥ %%y) := pure (y, x)
| _ := none

meta def match_le : expr → option (expr × expr)
| `(%%x ≤ %%y) := pure (x, y)
| `(%%x ≥ %%y) := pure (y, x)
| _ := none

meta def match_lt : expr → option (expr × expr)
| `(%%x < %%y) := pure (x, y)
| `(%%x > %%y) := pure (y, x)
| _ := none

#print list.filter_map

inductive edge
| lt | le

instance : has_to_string edge :=
⟨ λ e, match e with
      | edge.lt := "lt"
      | edge.le := "le"
      end ⟩

meta instance edge.has_to_format : has_to_format edge :=
⟨ λ e, to_fmt $ to_string e ⟩

meta instance rb_lmap.has_to_format {α β} [has_to_tactic_format α] [has_to_tactic_format β] : has_to_tactic_format (native.rb_lmap α β) :=
by delta native.rb_lmap; apply_instance

open native

meta def graph := (native.rb_lmap expr (expr × edge × expr))

meta instance graph.has_to_format : has_to_tactic_format graph :=
by delta graph; apply_instance

meta def dfs_trans' (g : graph) (r : ref expr_set) (v : expr) : edge → expr → expr → tactic expr
| e x h := do
  x ← instantiate_mvars x,
  -- trace!"visit {x}, going to {v}",
  vs ← read_ref r,
  -- trace!"seen: {vs}",
  if vs.contains x then failed
  else if v = x then pure h
  else do
    write_ref r $ vs.insert x,
    -- trace (g.find x),
    (g.find x).mfirst $ λ ⟨h',e',y⟩, do
      -- trace!"try: {x}, {y}",
      (e,h) ← match e, e' with
              | edge.lt, edge.lt := prod.mk edge.lt <$> mk_app ``lt_trans [h, h']
              | edge.lt, edge.le := prod.mk edge.lt <$> mk_app ``lt_of_lt_of_le [h, h']
              | edge.le, edge.lt := prod.mk edge.lt <$> mk_app ``lt_of_le_of_lt [h, h']
              | edge.le, edge.le := prod.mk edge.le <$> mk_app ``le_trans [h, h']
              end,
      -- trace"ok",
      dfs_trans' e y h

meta def dfs_trans (g : graph) (v v' : expr) : tactic expr :=
using_new_ref mk_expr_set $ λ r, do
  h ← mk_mapp ``le_refl [none, none, v],
  dfs_trans' g r v' edge.le v h

#check cc_state

lemma lt_of_eq_of_lt_of_eq {α} {R : α → α → Prop} {x x' y' y : α} (h₀ : x = x') (h₁ : R x' y') (h₂ : y' = y) :
  R x y := by subst_vars; exact h₁

-- lemma t_of_eq_of_lt_of_eq {α} [has_lt α] {x x' y' y : α} (h₀ : x = x') (h₁ : x' < y') (h₂ : y' = y) :
--   x < y := by subst_vars; exact h₁

meta def chain_trans : tactic unit := do
tgt ← target,
(x, y) ← match_le_or_lt tgt,
α ← infer_type x,
s ← cc_state.mk_using_hs,
ls ← local_context >>= list.mfilter_map'
  (λ h, do t ← infer_type h,
           do { (e, x, y) ← prod.mk edge.le <$> match_le t <|>
                    prod.mk edge.lt <$> match_lt t,
                let x' := s.root x,
                let y' := s.root y,
                x_pr ← s.eqv_proof x' x,
                y_pr ← s.eqv_proof y y',
                h' ← mk_app ``lt_of_eq_of_lt_of_eq [x_pr, h, y_pr],
                -- trace!"h  : {infer_type h}",
                -- trace!"h' : {infer_type h'}",
                infer_type x >>= is_def_eq α,
                pure [(h', e, x', y')] }),
 -- <|>
 --           do { (x, y) ← match_eq t,
 --                infer_type x >>= is_def_eq α,
 --                h₀ ← mk_eq_symm h >>= mk_app ``le_of_eq ∘ list.ret,
 --                h₁ ← mk_app ``le_of_eq [h],
 --                pure [(h₀, edge.le, y, x), (h₁, edge.le, x, y)] }),
let m := list.foldl  (λ (m : graph) (e : _ × _ × _ × _),
  let ⟨pr,e,x,y⟩ := e in
  m.insert x (pr,e,y)) (native.rb_lmap.mk expr (expr × edge × expr)) ls.join,
x ← whnf x,
y ← whnf y,
pr ← dfs_trans m x y,
tactic.apply pr <|>
  mk_app ``le_of_lt [pr] >>= tactic.apply,
skip

end interactive

setup_tactic_parser
precedence `=?`:0

import_private set_cases_tags
import_private cases_postprocess

meta def interactive.trichotomy (x : parse texpr) (_ : parse $ tk "=?") (y : parse texpr) (hyp : parse $ tk "with" *> ident <|> pure `h) : tactic unit := do
x ← to_expr x,
y ← to_expr y,
α ← infer_type x,
inst ← mk_app ``linear_order [α] >>= mk_instance,
h' ← mk_mapp ``cmp_compares [α, inst, x, y] >>= note hyp none,

e ← mk_app ``cmp [x, y],
n ← revert_kdependencies e,
x ← get_unused_name,
tactic.generalize e x,

h ← tactic.intro1,
s ← simp_lemmas.mk.add_simp ``ordering.compares,
in_tag ← get_main_tag,
focus1 $ do
  hs ← cases_core h,
  set_cases_tags in_tag $ cases_postprocess hs,
  gs ← get_goals,
  gs ← (gs.zip hs).mmap $ λ ⟨g,h,_,σ⟩, do
  { set_goals [g],
    interactive.propagate_tags $ do
    { dsimp_target (some s) [] { fail_if_unchanged := ff },
      intron_no_renames (n) },
    get_goals },
  set_goals gs.join

end tactic

example {x y : ℤ} (h : x ≤ y) (h : x < y) : true :=
begin
  find_hyp hh := x ≤ _ then { replace hh : x = y, admit },
  trivial
end

example {x y z w u v : ℕ} {a b : ℤ} (h₀ : x ≤ y) (h₁ : y < z) (h : y < u) (h₂ : z < w) (h₃ : w = u) (h₄ : u < v) : x ≤ u :=
begin
  chain_trans,
end


import data.finset.basic
import data.finset.sort
import data.finset.fold
import data.finmap
import data.fintype.basic
import tactic.linarith

/-
 - Model of Golang.
 - This version explores a strategy for putting the full field information
 - of a struct within its own type declaration. This would allow one to ensure
 - the struct is correct (and initiate instances of it) without refering to a
 - global context.
 - The downside to this approach is that you cannot have self-referential
 - structs, e.g. Linked List.
 -/

  inductive TypeDecl: Type
   | gInt:  TypeDecl
   | gStr:  TypeDecl
   | gErr:  TypeDecl
   | gBool: TypeDecl
   | gPtr:  TypeDecl → TypeDecl
   | gArr: ℕ → TypeDecl → TypeDecl
   | gStruct  (name: string) (fields: finset string)
              (vals: ∀ s ∈ fields, TypeDecl) : TypeDecl

  instance : decidable_eq TypeDecl :=
    begin
    intros a b, induction a generalizing b; cases b;
    try {{ apply decidable.is_false, rintro ⟨⟩ }},
    { simp [], exact decidable.true  },
    { simp [], exact decidable.true },
    { simp [], exact decidable.true },
    { simp [], exact decidable.true },
    {simp, exact (a_ih b), },
    { simp, apply @and.decidable (a_a=b_a) (a_a_1= b_a_1)
          (nat.decidable_eq a_a b_a) (a_ih b_a_1)},
    {refine if h: a_name = b_name then _ else decidable.is_false (by simp [h]),
     refine if h': a_fields = b_fields then _ else decidable.is_false (by simp [h']),
     subst h, subst h', simp, resetI,
     refine decidable_of_iff (∀ x h, (a_vals x h = b_vals x h))
         (by {intros, split,
            {intros,  simp [function.funext_iff],exact a,},
            {intros, simp [function.funext_iff] at a, exact a x h} }),},
    end

  open TypeDecl
  instance: inhabited TypeDecl := ⟨gInt⟩

  inductive PrimType: TypeDecl → Type
   | pInt:  ℤ → PrimType gInt
   | pStr:  string → PrimType gStr
   | pErr:  string → PrimType gErr
   | pBool: bool → PrimType gBool
   | pPtr (t: TypeDecl): ℤ → PrimType (gPtr t)
   | pArr (t: TypeDecl) (n:ℕ): (∀ s: fin n, PrimType t)
                               → PrimType (gArr n t)
   | pStruct  (name: string) (fields: finset string)
              (valtypes: ∀ s ∈ fields, TypeDecl)
              (vals: ∀ s ∈ fields, PrimType (valtypes s H))
              : PrimType (@gStruct name fields valtypes)
  open PrimType
  instance : ∀ t, decidable_eq (PrimType t) :=
    begin
    intros t a b, induction a generalizing b; cases b;
    try {{ apply decidable.is_false, rintro ⟨⟩ }},
    { exact decidable_of_iff (a = b_1) (by simp) },
    { exact decidable_of_iff (a = b_1) (by simp) },
    { exact decidable_of_iff (a = b_1) (by simp) },
    { exact decidable_of_iff (a = b_1) (by simp) },
    { exact decidable_of_iff (a_a = b_a) (by simp) },
    {resetI,
       refine decidable_of_iff (∀ x , a_a x = b_a x) ( by {intros, split,
        {intros, simp [function.funext_iff], exact a},
        {intros, simp [function.funext_iff] at a, exact a x}})},
    { resetI, refine decidable_of_iff (∀ x h, a_vals x h = b_vals x h)
        (by {intros, split,
        {intros, simp [function.funext_iff], exact a},
        {intros, simp [function.funext_iff] at a, exact a x h}}),},
    end

  def inh: ∀ t, (PrimType t)
   | gInt := pInt 0
   | gStr := pStr ""
   | gErr := pErr ""
   | gBool := pBool ff
   | (gPtr x) := @pPtr x 0
   | (gArr n t) := @pArr t n (λ _, inh t)
   | (@gStruct  name fs vs) := @pStruct _ _ _
      (λ x h, (have H: (vs x h).sizeof <
                        1 + name.length + finset.sizeof string fs := sorry,
      inh (vs x h)))

  instance: ∀ t, inhabited (PrimType t):= λ t, ⟨inh t⟩

  @[derive decidable_eq]
  structure Ptr (t: TypeDecl) := (val: ℤ)
  @[derive decidable_eq]
  structure Arr (n: ℕ) (t: TypeDecl) := (vals: (∀ s: fin n, PrimType t))

  structure Struct (name: string) (fields: finset string)
                         (valtypes: ∀ s ∈ fields, TypeDecl)
                          :=
    (vals: ∀ s ∈ fields, PrimType (valtypes s H))

lemma strext {name} {fields} {valtypes} (a b: Struct  name fields valtypes)
                (h : a.vals=b.vals) : a = b := begin
                  cases a, cases b, simp at h, simp, exact h
        end

  def type_dict: TypeDecl → Type
      | gStr    := string
      | gErr    := string
      | gInt    := ℤ
      | gBool   := bool
      | (gPtr t):= Ptr t
      | (gArr n t):= Arr n t
      | (gStruct name fs vs):= Struct name fs vs

  @[simp]
  def to_td: ∀ t, (PrimType t) → (type_dict t)
   | gInt            (pInt x)           := x
   | gStr            (pStr x)           := x
   | gErr            (pErr x)           := x
   | gBool           (pBool x)          := x
   | (gPtr _)        (pPtr _ x)         := Ptr.mk x
   | (gArr _ t)      (pArr _ _ xs)      := Arr.mk xs
   | (gStruct _ _ _) (pStruct _ _ _ xs) :=  ⟨xs⟩

  @[simp]
  def from_td: ∀ t, (type_dict t) → (PrimType t)
   | gInt x := pInt x
   | gStr x := pStr x
   | gErr x := pErr x
   | gBool x := pBool x
   | (gPtr t) x := pPtr t x.val
   | (gArr n t) x := pArr t n x.vals
   | (gStruct n fs vs) x := pStruct n fs vs x.vals

  /-
   - Bijection between the PrimType representation and Lean types
   -/
  def td_bij: ∀ t:TypeDecl, equiv (PrimType t) (type_dict t) := λ t,
  { to_fun    := to_td t,
    inv_fun   := from_td t,
    left_inv  := by {rintro ⟨n⟩; refl},
    right_inv := by {induction t;
      {unfold function.right_inverse,unfold function.left_inverse, intros x,
       cases x; simp  },
    }}

  def PrimType.to_typedecl  {t} (x: PrimType t): TypeDecl := t

  /-
   - Sometimes we don't know ahead of time if the type β that we're expecting
   - is actually equal to the type 'type_dict x' we are giving, so these as_type
   - functions handle this in a safe way.
   -/
  @[simp]
  def as_type  {dt: TypeDecl}
              (t: type_dict dt) (dt':TypeDecl): option (type_dict dt') :=
    if h: dt = dt' then some (by rw ← h; exact t) else none

  @[simp]
  def as_type_opt  {dt: TypeDecl}  (t: option (type_dict dt))
                  (dt': TypeDecl): option (type_dict dt') :=
      t >>= λ somet, @as_type dt somet dt'

section language_operators

  /-
   - Defining operators in Go
   -/
  @[derive decidable_eq]
  inductive Relop: Type | EQ | NEQ | LE

  @[derive decidable_eq]
  inductive Binop: Type | PLUS | MINUS | OR | AND

  @[derive decidable_eq]
  inductive Unop: Type | NOT | NEG

end language_operators

section gostructs
  open PrimType open nat

  lemma ltsucc : ∀ n, n < n+1 :=
    by simp only [nat.succ_pos', forall_const, lt_add_iff_pos_right]

  def to_fin (n) (i:ℕ) : option (fin n) :=
    if h: i < n then some ⟨i,h⟩ else none

  def shrink' (f:finset string) (h1: f.nonempty) (ts : ∀ s ∈ f, TypeDecl):
        ∀ s ∈ (f.erase (f.min' h1)), TypeDecl := by {
          intros s h,
          rw finset.mem_erase at h,
          exact ts s h.2}

  /-
   - Some analogue to structural recursion for linear ordered finsets
   -/
  def shrink (f : finset string) (ts : ∀ s ∈ f, TypeDecl)
              (vals: ∀ s ∈ f, PrimType (ts s H))
              : (Σ (f': finset string) (ts': ∀ s ∈ f', TypeDecl),
                  ∀ s ∈ f', PrimType (ts' s H)) :=
    if h0: f = ∅ then ⟨f, ts, λ a b, inh (ts a b)⟩
    else
      have h1: f.nonempty := finset.nonempty_of_ne_empty h0,
      have h2: ∃ a, a ∈ f.min := finset.min_of_nonempty h1,
      have mxmem: f.min' h1 ∈ f := finset.min'_mem f h1,
      have mxtyp: TypeDecl := ts (f.min' h1) mxmem,

      have newvals: ∀ s ∈ (f.erase (f.min' h1)), PrimType ((shrink' f h1 ts) s H), by {
          intros s h',
          rw finset.mem_erase at h', dedup,
        exact vals s h'_1.2},

      ⟨f.erase (f.min' h1), shrink' f h1 ts, newvals⟩
  /-
   - Add a field, value pair to an existing Struct (helper for add_key)
   -/
    def add_key' (k: string) (t: TypeDecl) (val: PrimType t)
                (f: finset string) (ts: ∀ s ∈ f, TypeDecl)
                (v : ∀ s ∈ f, PrimType (ts s H))
                : Σ(fields: finset string) (typs:∀ s ∈ fields, TypeDecl ),
                  ∀ s ∈ fields, PrimType (typs s H) :=
      ⟨insert k f, λ s mem, if h: s = k then t else  begin
        apply ts s, -- changes goal to needing to prove s is in old fields
          rw finset.mem_insert at mem, -- mem of insert means s=k or s ∈ f
            simp only [h, false_or] at mem, -- rule out poss. of s=k with ¬h
            exact mem end
        , λ s mem,
      if h: s=k then by {subst h, simp, exact val} else  begin
        rw finset.mem_insert at mem, -- mem of insert means s=k or s ∈ f
        simp only [h, false_or] at mem, -- rule out poss. of s=k with ¬h
          have Q : (dite (s = k) (λ (h : s = k), t) (λ (h : ¬s = k), ts s mem)
                    = ts s _), by { simp [h]},
        rw Q, exact v s mem
       end⟩

  def add_key: Π (s) (t), PrimType t → string
              → (Σ f v, Struct s f v) → (Σ f v, Struct s f v)
      | s t pt key ⟨f', ts', ⟨vs'⟩⟩  :=
    match add_key' key t pt f' ts' vs' with
      | ⟨f, ts, vs⟩ := ⟨f, ts, ⟨vs⟩⟩
     end

  /-
   - Convenience constructor for GoStruct
   - An arbitrary value can be used for the image of a map with an empty domain
   -/
  def GoStruct.from_list: Π(s) , list (string × (Σ t, PrimType t))
                          → Σ (f) (v), Struct s f v
   | s  []        := ⟨∅, (λ _ _, gBool), ⟨λ _ _, pBool ff⟩⟩ -- impossible
   | n ((k,⟨t,v⟩)::tl) := add_key n t v k (GoStruct.from_list n tl)

end gostructs

section to_strings

open nat
  def str_typeddecl: TypeDecl → string
    | gInt := "Int"
    | gStr := "Str"
    | gErr := "Error"
    | gBool := "Bool"
    | (gPtr t) := "*" ++ str_typeddecl t
    | (gArr n t) := "["++to_string n ++"]" ++ str_typeddecl t
    | (gStruct s _ _) := s
  instance: has_to_string TypeDecl := ⟨str_typeddecl⟩
  instance: has_repr      TypeDecl := ⟨str_typeddecl⟩

  /-
  - Rendering GoPrimTypes works but relies on sorry in two places to show that
  - it terminates. The struct's fields are printed in alphabetical order.
  -/
  meta mutual def str_ptype, str_arr, str_fields
  with str_ptype: ∀ t, PrimType t → string
   | _ (pInt p)            := to_string p
   | _ (pStr p)            := to_string p
   | _ (pErr p)            := to_string p
   | _ (pBool p)           := to_string p
   | _ (pPtr t v)            := "*" ++ to_string t ++ "@" ++ to_string v
   | (gArr n t) (pArr _ _ v) := "[" ++ str_arr n v ++ "]"
   | _ (@pStruct name f t v) := name ++"{"++ str_fields f t v ++"}"

  with str_arr: Π{n} {t}, ℕ → (fin n → PrimType t) → string
   | _ _ 0 _ := ""
   | n t (succ i) a := (to_fin n i).elim "" (λ x, str_ptype t $ a x)
                        ++ str_arr i a
 with str_fields: Π(f:finset string) (ts: (∀ s ∈ f, TypeDecl)),
               (∀ s ∈ f, PrimType (ts s H))  → string
  | f ts v :=
      begin
      have e: decidable (f = ∅) := finset.has_decidable_eq f ∅,
      cases e with hf ht,
          {have h1: f.nonempty := finset.nonempty_of_ne_empty hf,
           have head: string := f.min' h1 ++ ": " ++ (
                    str_ptype (ts (f.min' h1) (finset.min'_mem f h1))
                              (v (f.min' h1) (finset.min'_mem f h1)))
                  ++ if f.card > 1 then ", " else "",
            exact head ++ (
                match shrink f ts v with
                 | ⟨a,b,c⟩ := str_fields a b c
                end )},
          {exact ""}
        end

  meta def str_gstruct (s) (f) (v) (x: Struct s f v): string :=
    str_ptype _ (pStruct s f v x.vals)

  meta instance :∀ t,  has_to_string (PrimType t):= λ t, ⟨str_ptype t⟩
  meta instance :∀ t,  has_to_string (Ptr t):=
    λ t, ⟨λ ⟨z⟩, str_ptype (gPtr t) (pPtr t z)⟩
  meta instance :∀ n t,  has_to_string (Arr n t):=
    λ n t, ⟨λ ⟨z⟩, str_ptype (gArr n t) (pArr t n z)⟩

  meta instance: Π(s) (f) (v), has_repr (Struct s f v)
    | s f v := ⟨str_gstruct s f v⟩.
  meta instance: Π(s) (f) (v), has_to_string (Struct s f v)
    | s f v := ⟨str_gstruct s f v⟩

 meta instance type_dict_to_string: Π(x:TypeDecl), has_to_string (type_dict x)
    | gStr     := string.has_to_string
    | gErr     := string.has_to_string
    | gInt     := int.has_to_string
    | gBool    := bool.has_to_string
    | (gPtr t) := (Ptr.has_to_string t)
    | (gArr n t) := (Arr.has_to_string n t)
    | (gStruct n f v) := (Struct.has_to_string n f v)

end to_strings

section other_instances

  open TypeDecl

  instance int_le:     has_le     (type_dict gInt) := int.has_le
  instance int_add:    has_add    (type_dict gInt) := int.has_add
  instance int_lt:     has_lt     (type_dict gInt) := int.has_lt
  instance int_zero:   has_zero   (type_dict gInt) := int.has_zero
  instance int_one:    has_one    (type_dict gInt) := int.has_one
  instance int_neg:    has_neg    (type_dict gInt) := int.has_neg
  instance int_sub:    has_sub    (type_dict gInt) := int.has_sub
  instance str_le:     has_le     (type_dict gStr) := string.has_le
  instance str_append: has_append (type_dict gStr) := string.has_append

end other_instances

/-

--   /-
--    - Two lemmas about fieldlist which help prove define GoStruct.index
--    -/
--   lemma fl_len_eq {s} (x: GoStruct s): x.fieldlist.length = x.len := begin
--     induction x,
--     simp, unfold GoStruct.fieldlist, simp,
--   end

--   lemma fl_mem {s} {x: GoStruct s} {key: string}:
--                key ∈ x.fieldlist → key ∈ x.fields  := begin
--     intros H, induction x,
--     unfold GoStruct.fieldlist at H, simp at H, simp, exact H
--   end

--   /-
--    - Alternative of nth_le that also returns a proof that the return element is a
--    - member of the original list being indexed.
--    -/
--   @[simp]
--   def nth_le_mem : Π {α: Type} (l : list α) (n),
--                       n < l.length → psigma (λ a:α, a ∈ l)
--    | _ []       n     h := absurd h (nat.not_lt_zero n)
--    | _ (a :: l) 0     h := ⟨a, by {simp}⟩
--    | _ (a :: l) (n+1) h :=
--       match (nth_le_mem l n (nat.le_of_succ_le_succ h)) with
--             | ⟨q, r⟩ := ⟨q, list.mem_cons_of_mem a r⟩
--       end

--   lemma ltsucc : ∀ n, n < n+1 :=
--     by simp only [nat.succ_pos', forall_const, lt_add_iff_pos_right]

--   /-
--    - Given an index into the (sorted) fields, return a string and a proof that it
--    - is a member of the finset.
--    -/
--   def GoStruct.index {name} {n:ℕ} (x: GoStruct name)
--                      (H: n < x.len) : psigma (λ f, f ∈ x.fields) := begin
--     have z := nth_le_mem x.fieldlist n
--                                     (by {rw fl_len_eq x, exact H}),
--     exact ⟨z.1, fl_mem z.2⟩
--   end

--   /-
--    - Helper function from GoStruct.structvals
--    -/
--   def structvals': Π {n:ℕ} {s: string} {x:GoStruct s},
--                     n < x.len → list (string × GoPrimType)
--     |  0 _  _ _  := []
--     |  (n+1) _ x h := match x.index h with
--                           | ⟨u,v ⟩ :=  (u, x.vals u v) ::
--                                       structvals' (lt.trans (ltsucc n) h)
--                       end

--   /-
--    - key sorted list of values
--    -/
--   def GoStruct.pairs {s} (x: GoStruct s) : list (string × GoPrimType) :=
--     begin
--       cases h : x.len with n,
--         {exact []},
--       {have H: x.len -1 < x.len, by {rw h, simp, exact ltsucc n},
--         exact structvals' H}
--   end

--   /-
--    - Access a field of a struct instance with an expected type
--    -/
--   def check_member_fields : Π(α: GoTypeDecl) {s: string},
--       option (GoStruct s) → string → option (type_dict α)
--     | _ _ none                          _     := none
--     | α _ (some s) field := GoStruct.get s field >>=
--                               λx, as_type (x.to_gotype.2) α

-/


/-
Copyright (c) 2017 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Gabriel Ebner

! This file was ported from Lean 3 source module data.buffer.parser
! leanprover-community/mathlib commit 549e2fed50b361d0d49a3dd1e7ccb6de9440059b
! Please do not edit these lines, except to modify the commit id
! if you have ported upstream changes.
-/
import Leanbin.Data.Buffer
import Leanbin.Data.Dlist

inductive ParseResult (α : Type)
  | done (pos : ℕ) (result : α) : ParseResult
  | fail (pos : ℕ) (expected : Dlist String) : ParseResult
#align parse_result ParseResult

/--
The parser monad. If you are familiar with the Parsec library in Haskell, you will understand this.  -/
def Parser (α : Type) :=
  ∀ (input : CharBuffer) (start : ℕ), ParseResult α
#align parser Parser

namespace Parser

variable {α β γ : Type}

protected def bind (p : Parser α) (f : α → Parser β) : Parser β := fun input pos =>
  match p input Pos with
  | ParseResult.done Pos a => f a input Pos
  | ParseResult.fail Pos expected => ParseResult.fail Pos expected
#align parser.bind Parser.bind

protected def pure (a : α) : Parser α := fun input pos => ParseResult.done Pos a
#align parser.pure Parser.pure

private theorem parser.id_map (p : Parser α) : Parser.bind p Parser.pure = p :=
  by
  apply funext; intro input
  apply funext; intro pos
  dsimp only [Parser.bind]
  cases p input Pos <;> exact rfl
#align parser.parser.id_map parser.parser.id_map

private theorem parser.bind_assoc (p : Parser α) (q : α → Parser β) (r : β → Parser γ) :
    Parser.bind (Parser.bind p q) r = Parser.bind p fun a => Parser.bind (q a) r :=
  by
  apply funext; intro input
  apply funext; intro pos
  dsimp only [Parser.bind]
  cases p input Pos <;> try dsimp only [bind]
  cases q result input pos_1 <;> try dsimp only [bind]
  all_goals rfl
#align parser.parser.bind_assoc parser.parser.bind_assoc

protected def fail (msg : String) : Parser α := fun _ pos =>
  ParseResult.fail Pos (Dlist.singleton msg)
#align parser.fail Parser.fail

instance : Monad Parser where
  pure := @Parser.pure
  bind := @Parser.bind

instance : LawfulMonad Parser where
  id_map := @Parser.id_map
  pure_bind _ _ _ _ := rfl
  bind_assoc := @Parser.bind_assoc

instance : MonadFail Parser :=
  { Parser.monad with fail := @Parser.fail }

protected def failure : Parser α := fun _ pos => ParseResult.fail Pos Dlist.empty
#align parser.failure Parser.failure

protected def orelse (p q : Parser α) : Parser α := fun input pos =>
  match p input Pos with
  | ParseResult.fail pos₁ expected₁ =>
    if pos₁ ≠ Pos then ParseResult.fail pos₁ expected₁
    else
      match q input Pos with
      | ParseResult.fail pos₂ expected₂ =>
        if pos₁ < pos₂ then ParseResult.fail pos₁ expected₁
        else
          if pos₂ < pos₁ then ParseResult.fail pos₂ expected₂
          else-- pos₁ = pos₂
              ParseResult.fail
              pos₁ (expected₁ ++ expected₂)
      | ok => ok
  | ok => ok
#align parser.orelse Parser.orelse

instance : Alternative Parser where
  failure := @Parser.failure
  orelse := @Parser.orelse

instance : Inhabited (Parser α) :=
  ⟨Parser.failure⟩

/-- Overrides the expected token name, and does not consume input on failure. -/
def decorateErrors (msgs : Thunk (List String)) (p : Parser α) : Parser α := fun input pos =>
  match p input Pos with
  | ParseResult.fail _ expected => ParseResult.fail Pos (Std.DList.lazy_ofList (msgs ()))
  | ok => ok
#align parser.decorate_errors Parser.decorateErrors

/-- Overrides the expected token name, and does not consume input on failure. -/
def decorateError (msg : Thunk String) (p : Parser α) : Parser α :=
  decorateErrors [msg ()] p
#align parser.decorate_error Parser.decorateError

/-- Matches a single character. Fails only if there is no more input. -/
def anyChar : Parser Char := fun input pos =>
  if h : Pos < input.size then
    let c := input.read ⟨Pos, h⟩
    ParseResult.done (Pos + 1) c
  else ParseResult.fail Pos Dlist.empty
#align parser.any_char Parser.anyChar

/-- Matches a single character satisfying the given predicate. -/
def sat (p : Char → Prop) [DecidablePred p] : Parser Char := fun input pos =>
  if h : Pos < input.size then
    let c := input.read ⟨Pos, h⟩
    if p c then ParseResult.done (Pos + 1) c else ParseResult.fail Pos Dlist.empty
  else ParseResult.fail Pos Dlist.empty
#align parser.sat Parser.sat

/-- Matches the empty word. -/
def eps : Parser Unit :=
  return ()
#align parser.eps Parser.eps

/-- Matches the given character. -/
def ch (c : Char) : Parser Unit :=
  decorateError c.toString <| sat (· = c) >> eps
#align parser.ch Parser.ch

/-- Matches a whole char_buffer.  Does not consume input in case of failure. -/
def charBuf (s : CharBuffer) : Parser Unit :=
  decorateError s.toString <| s.toList.mapM' ch
#align parser.char_buf Parser.charBuf

/-- Matches one out of a list of characters. -/
def oneOf (cs : List Char) : Parser Char :=
  (decorateErrors do
      let c ← cs
      return c) <|
    sat (· ∈ cs)
#align parser.one_of Parser.oneOf

def oneOf' (cs : List Char) : Parser Unit :=
  oneOf cs >> eps
#align parser.one_of' Parser.oneOf'

/-- Matches a string.  Does not consume input in case of failure. -/
def str (s : String) : Parser Unit :=
  decorateError s <| s.toList.mapM' ch
#align parser.str Parser.str

/-- Number of remaining input characters. -/
def remaining : Parser ℕ := fun input pos => ParseResult.done Pos (input.size - Pos)
#align parser.remaining Parser.remaining

/-- Matches the end of the input. -/
def eof : Parser Unit :=
  decorateError "<end-of-file>" do
    let rem ← remaining
    guard <| rem = 0
#align parser.eof Parser.eof

def foldrCore (f : α → β → β) (p : Parser α) (b : β) : ∀ reps : ℕ, Parser β
  | 0 => failure
  | reps + 1 =>
    (do
        let x ← p
        let xs ← foldr_core reps
        return (f x xs)) <|>
      return b
#align parser.foldr_core Parser.foldrCore

/-- Matches zero or more occurrences of `p`, and folds the result. -/
def foldr (f : α → β → β) (p : Parser α) (b : β) : Parser β := fun input pos =>
  foldrCore f p b (input.size - Pos + 1) input Pos
#align parser.foldr Parser.foldr

def foldlCore (f : α → β → α) : ∀ (a : α) (p : Parser β) (reps : ℕ), Parser α
  | a, p, 0 => failure
  | a, p, reps + 1 =>
    (do
        let x ← p
        foldl_core (f a x) p reps) <|>
      return a
#align parser.foldl_core Parser.foldlCore

/-- Matches zero or more occurrences of `p`, and folds the result. -/
def foldl (f : α → β → α) (a : α) (p : Parser β) : Parser α := fun input pos =>
  foldlCore f a p (input.size - Pos + 1) input Pos
#align parser.foldl Parser.foldl

/-- Matches zero or more occurrences of `p`. -/
def many (p : Parser α) : Parser (List α) :=
  foldr List.cons p []
#align parser.many Parser.many

def manyChar (p : Parser Char) : Parser String :=
  List.asString <$> many p
#align parser.many_char Parser.manyChar

/-- Matches zero or more occurrences of `p`. -/
def many' (p : Parser α) : Parser Unit :=
  many p >> eps
#align parser.many' Parser.many'

/-- Matches one or more occurrences of `p`. -/
def many1 (p : Parser α) : Parser (List α) :=
  List.cons <$> p <*> many p
#align parser.many1 Parser.many1

/-- Matches one or more occurences of the char parser `p` and implodes them into a string. -/
def manyChar1 (p : Parser Char) : Parser String :=
  List.asString <$> many1 p
#align parser.many_char1 Parser.manyChar1

/-- Matches one or more occurrences of `p`, separated by `sep`. -/
def sepBy1 (sep : Parser Unit) (p : Parser α) : Parser (List α) :=
  List.cons <$> p <*> many (sep >> p)
#align parser.sep_by1 Parser.sepBy1

/-- Matches zero or more occurrences of `p`, separated by `sep`. -/
def sepBy (sep : Parser Unit) (p : Parser α) : Parser (List α) :=
  sepBy1 sep p <|> return []
#align parser.sep_by Parser.sepBy

def fixCore (F : Parser α → Parser α) : ∀ max_depth : ℕ, Parser α
  | 0 => failure
  | max_depth + 1 => F (fix_core max_depth)
#align parser.fix_core Parser.fixCore

/-- Matches a digit (0-9). -/
def digit : Parser Nat :=
  decorateError "<digit>" do
    let c ← sat fun c => '0' ≤ c ∧ c ≤ '9'
    pure <| c - '0'.toNat
#align parser.digit Parser.digit

/-- Matches a natural number. Large numbers may cause performance issues, so
don't run this parser on untrusted input. -/
def nat : Parser Nat :=
  decorateError "<natural>" do
    let digits ← many1 digit
    pure <|
        Prod.fst <|
          digits (fun digit ⟨Sum, magnitude⟩ => ⟨Sum + digit * magnitude, magnitude * 10⟩) ⟨0, 1⟩
#align parser.nat Parser.nat

/-- Fixpoint combinator satisfying `fix F = F (fix F)`. -/
def fix (F : Parser α → Parser α) : Parser α := fun input pos =>
  fixCore F (input.size - Pos + 1) input Pos
#align parser.fix Parser.fix

private def make_monospaced : Char → Char
  | '\n' => ' '
  | '\t' => ' '
  | '\x0d' => ' '
  | c => c
#align parser.make_monospaced parser.make_monospaced

def mkErrorMsg (input : CharBuffer) (pos : ℕ) (expected : Dlist String) : CharBuffer :=
  let left_ctx := (input.take Pos).takeRight 10
  let right_ctx := (input.drop Pos).take 10
  (left_ctx.map makeMonospaced ++ right_ctx.map makeMonospaced ++ "\n".toCharBuffer ++
              left_ctx.map fun _ => ' ') ++
            "^\n".toCharBuffer ++
          "\n".toCharBuffer ++
        "expected: ".toCharBuffer ++
      String.toCharBuffer (" | ".intercalate expected.toList) ++
    "\n".toCharBuffer
#align parser.mk_error_msg Parser.mkErrorMsg

/-- Runs a parser on the given input.  The parser needs to match the complete input. -/
def run (p : Parser α) (input : CharBuffer) : Sum String α :=
  match (p <* eof) input 0 with
  | ParseResult.done Pos res => Sum.inr res
  | ParseResult.fail Pos expected => Sum.inl <| Buffer.toString <| mkErrorMsg input Pos expected
#align parser.run Parser.run

/-- Runs a parser on the given input.  The parser needs to match the complete input. -/
def runString (p : Parser α) (input : String) : Sum String α :=
  run p input.toCharBuffer
#align parser.run_string Parser.runString

end Parser