import .util

meta def unwrap_lm_response (ident : option string) : json → tactic (list string)
| (json.array $ tactic_strings) := do
    result ← tactic_strings.mmap $ lift_option ∘ json.get_string,
    pure result
| exc := tactic.fail format!"{ident.get_or_else \"[unwrap_lm_response.anonymous]\"} run_best_beam_candidate UNEXPECTED: {exc}"

meta def json_float_array_sum : json → option json
| (json.array xs) := json.of_float <$> xs.mfoldr (λ msg acc, match msg with
  | (json.of_float val) := pure $ acc + val
  | (json.of_int val) := pure $ acc + native.float.of_int val
  | exc := none
  end) (0.0 : native.float)
| exc := none

-- WARNING(jesse, January 14 2021, 10:32 AM): instead of failing like `unwrap_lm_response`, this return an empty list when seeing an unexpected message
meta def unwrap_lm_response_logprobs (ident : option string) : json → tactic (list $ string × native.float)
| (json.array $ [(json.array predictions), (json.array scores)]) := do {
  decoded_strings ← predictions.mmap $ lift_option ∘ json.get_string,
  decoded_scores ← scores.mmap $ lift_option ∘ json.get_float,
  pure $ list.zip decoded_strings decoded_scores
}
| exc := tactic.trace format!"{ident.get_or_else \"[unwrap_lm_response_logprobs.anonymous]\"} run_best_beam_candidate UNEXPECTED: {exc}" *> pure []

section json

-- for debugging

meta def json.compare : Π (x y : json), bool
| (json.of_string s) (json.of_string s') := s = s'
| (json.of_int k) (json.of_int k') := k = k'
| (json.of_float x) (json.of_float x') := x = x' -- might have to make this tt
| (json.of_bool b) (json.of_bool b') := b = b'
| (json.null) (json.null) := tt
| (json.object kvs) (json.object kvs') := (list.zip kvs kvs').foldr
  (λ ⟨⟨k₁, v₁⟩, ⟨k₂, v₂⟩⟩ acc,
  json.compare k₁ k₂ && json.compare v₁ v₂ && acc) tt
| (json.array args) (json.array args') := (list.zip args args').foldr
  (λ ⟨j₁, j₂⟩ acc, acc && json.compare j₁ j₂) tt
| _ _ := ff

meta def json.to_raw_fmt : json → format
| (json.of_string s) := format!"(json.of_string \"{s}\")"
| (json.of_int k) := format!"(json.of_int {k})"
| (json.of_float x) := format!"(json.of_float {x})"
| (json.of_bool b) := format!"(json.of_bool {b})"
| (json.null) := "(json.null)"
| (json.object kvs) := let f : string × json → format :=
  (λ ⟨k,v⟩, json.to_raw_fmt k ++ " : " ++ json.to_raw_fmt v) in
   format!"(json.object " ++ format.join_using " " (f <$> kvs) ++ ")"
| (json.array args) := "(json.array " ++ format.join_using " " (json.to_raw_fmt <$> args) ++ ")"

end json

section derive_has_json
section has_to_tactic_json_name
open name

-- meta def has_to_tactic_json_name_aux : name → json
-- | anonymous := json.object $ [("constr", json.of_string "name.anonymous"), ("args", json.array [])]
-- | (mk_string str nm) := json.object $ [("constr", json.of_string "name.mk_string"),
--     ("args", json.array [json.of_string str, has_to_tactic_json_name_aux nm])]
-- | (mk_numeral u nm) := json.object $ [("constr", json.of_string "name.mk_numeral"),
--     ("args", json.array [json.of_int (int.of_nat ∘ unsigned.to_nat $ u), has_to_tactic_json_name_aux nm])]

meta def has_to_tactic_json_name_aux : name → json
| anonymous := json.array ∘ pure $ json.of_string "name.anonymous"
| (mk_string str nm) := json.array $
  [json.of_string "name.mk_string", json.of_string str, has_to_tactic_json_name_aux nm]
| (mk_numeral u nm) := json.array $
  [json.of_string "name.mk_numeral", json.of_int (int.of_nat u.to_nat), has_to_tactic_json_name_aux nm]

 -- json.object $ [("constr", json.of_string "name.mk_numeral"),
 --    ("args", json.array [json.of_int (int.of_nat ∘ unsigned.to_nat $ u), has_to_tactic_json_name_aux nm])]

meta instance : has_to_tactic_json name :=
⟨pure ∘ has_to_tactic_json_name_aux⟩

end has_to_tactic_json_name

section has_from_json_name

-- meta def has_from_json_name_aux : json → tactic name
-- | arg@(json.object [("constr", json.of_string c), ("args", json.array args)]) := do {
--   tactic.trace format!"GOT MATCH: {arg}",
--   match c with
--   | "name.anonymous" := pure name.anonymous
--   | "name.mk_string" := do {
--     (str, nm_json) ← (do {
--       [json.of_string str, nm_json] ← pure args,
--       pure (str, nm_json)
--     } <|> tactic.fail
--             format!"[has_from_json_name_aux.inner_match_1.mk_string] unexpected: {args}"),
--     name.mk_string str <$> has_from_json_name_aux nm_json
--   }
--   | "name.mk_numeral" := do {
--     (u, nm_json) ← (do {
--       [json.of_int u, nm_json] ← pure args,
--       pure (u, nm_json)
--     } <|> tactic.fail
--             format!"[has_from_json_name_aux.inner_match_1.mk_numeral] unexpected: {args}"),
--     name.mk_numeral (unsigned.of_nat ∘ int.to_nat $ u) <$> has_from_json_name_aux nm_json
--   }
--   | exc := tactic.fail format!"[has_from_json_name_aux.inner_match_1] unexpected: {exc}"
--   end
-- }
-- | exc := tactic.fail format!"[has_from_json_name_aux] unexpected: {exc}"

meta def has_from_json_name_aux : json → tactic name
| arg@(json.array (c::args)) := do {
  -- tactic.trace format!"GOT MATCH: {arg}",
  match c with
  | "name.anonymous" := pure name.anonymous
  | "name.mk_string" := do {
    (str, nm_json) ← (do {
      [json.of_string str, nm_json] ← pure args,
      pure (str, nm_json)
    } <|> tactic.fail
            format!"[has_from_json_name_aux.inner_match_1.mk_string] unexpected: {args}"),
    name.mk_string str <$> has_from_json_name_aux nm_json
  }
  | "name.mk_numeral" := do {
    (u, nm_json) ← (do {
      [json.of_int u, nm_json] ← pure args,
      pure (u, nm_json)
    } <|> tactic.fail
            format!"[has_from_json_name_aux.inner_match_1.mk_numeral] unexpected: {args}"),
    name.mk_numeral (unsigned.of_nat ∘ int.to_nat $ u) <$> has_from_json_name_aux nm_json
  }
  | exc := tactic.fail format!"[has_from_json_name_aux.inner_match_1] unexpected: {exc}"
  end
}
| exc := tactic.fail format!"[has_from_json_name_aux] unexpected: {exc}"

meta instance : has_from_json name :=
⟨has_from_json_name_aux⟩

end has_from_json_name

open tactic
namespace tactic
namespace interactive

meta def mk_to_tactic_json (type : name) : tactic unit := do {
  ls ← local_context,
  (x::_) ← tactic.intro_lst [`arg],
  et ← infer_type x,
  xs ← tactic.induction x,
  xs.mmap' $ λ ⟨c, args, _⟩, do
    (args', rec_call) ← args.mpartition $ λ e, do {e' ← tactic.to_expr ``(tactic json), bnot <$> e'.occurs <$> tactic.infer_type e},
    args'' ← args'.mmap (λ a, flip prod.mk a <$> (et.occurs <$> tactic.infer_type a)),
    let fn : list (bool × expr) → state_t (list expr) tactic (list expr) := λ args'', do {
      let pop : state_t (list expr) tactic (option expr) := do {
        xs ← get,
        match xs with
        | (a::as) := modify (λ _, as) *> pure (some a)
        | [] := pure none
        end
      },
      args''.mmap (λ ⟨b, a⟩, if b then do (some x) ← pop, pure x else state_t.lift $ do
      a_tp ← infer_type a,
      _inst ← mk_app ``has_to_tactic_json [a_tp] >>= mk_instance,
      tactic.to_expr ``(@has_to_tactic_json.to_tactic_json _ (%%_inst) %%a))
    },
    args''' ← prod.fst <$> (fn args'').run rec_call,

    c ← tactic.resolve_constant c,
    refine ``((λ (ys : list $ tactic json),
      (λ x, json.array [has_to_tactic_json_name_aux %%c,
        json.array x]) <$>  ys.mmap id) _), -- lol
    args'''.mmap (λ e, refine ``(list.cons %%e _)),
    tactic.to_expr ``(([] : list (tactic json))) >>= tactic.exact
}

meta def derive_has_to_tactic_json (pre : option name) : tactic unit := do {
  vs ← local_context,
  `(has_to_tactic_json %%f) ← target,
  env ← get_env,
  let n := f.get_app_fn.const_name,
  d ← get_decl n,
  refine ``( { to_tactic_json := _ } ),
  tgt ← target,
  extract_def (with_prefix pre n <.> "to_tactic_json") ff $ mk_to_tactic_json n
}

meta def has_to_tactic_json_derive_handler' (nspace : option name := none) : derive_handler :=
higher_order_derive_handler ``has_to_tactic_json (derive_has_to_tactic_json nspace) [] nspace

@[derive_handler]
meta def has_to_tactic_json_derive_handler : derive_handler :=
guard_class ``has_to_tactic_json has_to_tactic_json_derive_handler'

end interactive
end tactic
end derive_has_json

section derive_from_json
open tactic
namespace tactic

namespace interactive

meta def get_constr_and_args (arg : json) : option (string × list json) :=
match arg with
| (json.array [json.of_string c, json.array args]) := pure (c, args)
| _ := none
end

-- meta def json_to_expr : Π (arg : json), tactic expr
-- | (json.object [("constr", nm_json), ("args", json.array args)]) := do {
--   -- let c_nm := hacky_name_from_string c,
--   c_nm ← (json_to_expr nm_json) >>= eval_expr name,
--   constr ← mk_const c_nm,
--   if args.length = 0 then do {
--   tp ← tactic.infer_type constr,
--   e_id ← to_expr ``(@id %%tp),
--   pure $ e_id.mk_app [constr] -- ???????????????
--   }
--   else do
--   constr.mk_app <$> (args.mmap json_to_expr)
-- }
-- | arg@(json.of_int k) := do { -- WARNING: this is a hack for now
--   pure `(int.to_nat k)
-- }
-- | exc@(json.array $ x@(((json.of_string c))::rest)) := if ("name" < c) then do nm ← has_from_json_name_aux x, pure $ (by apply_instance : has_reflect name) nm else tactic.fail format!"[json_to_expr.name] unexpected: {exc}"
-- -- | arg@(json.of_bool b) := do { -- WARNING: this is a hack for now
-- --   pure `(tt)
-- -- }
-- -- -- TODO(jesse): as needed, built special built-in logic to handle constants
-- -- -- TODO(jesse): use `resolve_name` on `nm` to get the `has_from_json` instance
-- -- -- then make a recursive call with `json_to_expr`
-- -- -- better yet, move this logic into `mk_from_json`

-- -- | arg@(json.object [("builtin", nm_json), ("val", val_json)]) := do {
-- --   env ← get_env,
-- --   nm ← (has_from_json_name_aux nm_json),
-- --   -- d ← env.get nm,
-- --   ty_reflected ← declaration.type <$> get_decl nm,

-- --   _inst ← to_expr ``(has_from_json %%nm) >>= mk_instance,
-- --   _inst2 ← to_expr ``(has_reflect %%nm) >>= mk_instance,
-- --   to_expr ``(has_from_json.from_json _ %%_inst %%arg)
-- --   -- to_expr ``(has_from_json.from_json %%_inst $ %%arg)
-- --   }
-- | exc := tactic.fail format!"[json_to_expr] unexpected: {exc}"

meta def json_to_expr : Π (arg : json), tactic unit
| (json.array $ [nm_json, json.array args]) := do {
  -- let c_nm := hacky_name_from_string c,
  c_nm ← (has_from_json_name_aux nm_json),
  constr ← mk_const c_nm,
  if args.length = 0 then do {
  tp ← tactic.infer_type constr,
  e_id ← to_expr ``(@id %%tp),
  tactic.apply (e_id.mk_app [constr]) *> pure ()
  }
  else do
  -- tactic.apply constr.mk_app <$> (args.mmap json_to_expr)
  tactic.apply constr,
  args.mmap' json_to_expr
}
| arg@(json.of_int k) := do { -- WARNING: this is a hack for now
  tactic.exact `(int.to_nat k)
}
| exc@(json.array $ x@(((json.of_string c))::rest)) := if ("name" < c) then do nm ← has_from_json_name_aux x, tactic.exact ((by apply_instance : has_reflect name) nm) else tactic.fail format!"[json_to_expr.name] unexpected: {exc}"
-- | arg@(json.of_bool b) := do { -- WARNING: this is a hack for now
--   pure `(tt)
-- }
-- -- TODO(jesse): as needed, built special built-in logic to handle constants
-- -- TODO(jesse): use `resolve_name` on `nm` to get the `has_from_json` instance
-- -- then make a recursive call with `json_to_expr`
-- -- better yet, move this logic into `mk_from_json`

-- | arg@(json.object [("builtin", nm_json), ("val", val_json)]) := do {
--   env ← get_env,
--   nm ← (has_from_json_name_aux nm_json),
--   -- d ← env.get nm,
--   ty_reflected ← declaration.type <$> get_decl nm,

--   _inst ← to_expr ``(has_from_json %%nm) >>= mk_instance,
--   _inst2 ← to_expr ``(has_reflect %%nm) >>= mk_instance,
--   to_expr ``(has_from_json.from_json _ %%_inst %%arg)
--   -- to_expr ``(has_from_json.from_json %%_inst $ %%arg)
--   }
| exc := tactic.fail format!"[json_to_expr] unexpected: {exc}"

/- TODO(jesse): probably a better, recursive/lazy way of doing this -/
meta def mk_from_json (pre : option name) : tactic unit := do {
  (x::_) ← tactic.intro_lst [`_arg],
  real_tgt@`(tactic %%tgt) ← target,
  y ← to_expr ``(json_to_expr %%x),
  -- let real_tgt := ``(reflected $ tactic %%tgt),

  -- _ ← to_expr ``(do %%y >>= eval_expr unit),
  -- tactic.read >>= λ ts, tactic.trace format!"TACTIC STATE AFTER EVAL: {ts}",
  -- pure ()
  -- tac ← eval_expr (tactic unit) y,
  -- tac
  -- `(tactic %%tgt) ← target,
  -- -- y ← to_expr ``(@id (tactic %%tgt)) >>= (λ f, pure $ f.mk_app [y]),
  -- -- tactic.trace format!"OK?: {y}",
  -- -- tactic.unfreeze_local_instances,

  -- OK, i think this approach was the right one since `eval_expr` has type exactly `tactic alpha`
  -- although maybe we could also do... pure?
  -- before this line, y is a reflection of `json_to_expr %%x`, which returns an expr which is supposed to be the target
  -- to force this to evaluate, we turns this entire thing into a thing of type tactic tgt,

  -- result ← to_expr ``(do %%y >>= eval_expr %%tgt),
  gs ← tactic.get_goals,

  -- m ← tactic.mk_meta_var real_tgt,
  -- let m' := `([m]).to_expr,

  -- result ← to_expr ``(do %%y >>= eval_expr %%tgt),
  -- this doesn't work because `y` is just `json_to_expr %%x`, which contains the open variable `%%x`
  -- result ← to_expr ``(do %%y) >>= eval_expr expr,

  result ← to_expr ``(do try trivial, m ← mk_mvar, set_goals [m], %%y, tactic.get_assignment m >>= eval_expr %%tgt),
  tactic.exact result *> done
  -- (eval_expr json x) >>= json_to_expr
}

-- #check get_constr_and_args
-- meta def mk_from_json (pre : option name) : tactic unit := do {
--   (x::_) ← tactic.intro_lst [`_arg],
--   let rec := `(pure () : tactic unit).to_expr,
--   f ← to_expr ``(match (get_constr_and_args %%x) with
--   | (some ⟨constr, args⟩) := %%rec
--   | none := tactic.fail "[mk_from_json] unexpected failure"
--   end),
--   tactic.fail "NYI"

-- -- (constr, args))
-- }

meta def derive_has_from_json (pre : option name) : tactic unit := do {
  vs ← local_context,
  `(has_from_json %%f) ← target,
  env ← get_env,
  let n := f.get_app_fn.const_name,
  d ← get_decl n,
  refine ``( { from_json := _ } ),
  tgt ← target,
  let extract_def_nm := (with_prefix pre n <.> "from_json"),
  extract_def extract_def_nm ff $ mk_from_json n
}

meta def has_from_json_derive_handler' (nspace : option name := none) : derive_handler :=
higher_order_derive_handler ``has_from_json (derive_has_from_json nspace) [] nspace

@[derive_handler]
meta def has_from_json_derive_handler : derive_handler :=
guard_class ``has_from_json has_from_json_derive_handler'

end interactive

end tactic

end derive_from_json

section test

-- @[derive [has_to_tactic_json, has_from_json]]

inductive my_nat' : Type
-- /- `(x : α)` -/
-- | foo : DUH'
-- /- `{x : α}` -/
| bar : my_nat'
| baz : my_nat' → my_nat'
-- /- `⦃x:α⦄` -/
-- | strict_implicit : DUH'
-- /- `[x : α]`. Should be inferred with typeclass resolution. -/
-- | inst_implicit : DUH'
-- /- Auxiliary internal attribute used to mark local constants representing recursive functions
--         in recursive equations and `match` statements. -/
-- | aux_decl : DUH'

attribute [derive has_to_tactic_json] my_nat'
attribute [derive has_from_json] my_nat'

attribute [derive [has_reflect]] my_nat
attribute [derive [has_to_tactic_json]] my_nat
attribute [derive [has_from_json]] my_nat

-- run_cmd (has_to_tactic_json.to_tactic_json (my_nat.zero) >>= (has_from_json.from_json : json → tactic my_nat) >>= tactic.trace)

-- meta instance : has_to_format my_nat :=
-- ⟨λ x, match x with | my_nat.zero := "my_nat.zero" | (my_nat.succ x) := "my_nat.succ " ++ (by exact _match x) end⟩

-- meta instance : has_from_json my_nat :=
-- ⟨λ k,
-- -- match k with
-- -- | (json.object $ [("constr", c), ("args", (json.array args))]) := do
-- --   nm ← has_from_json_name_aux c,
-- --   match nm with
-- --   | `my_nat.zero := pure $ my_nat.zero
-- --   | `my_nat.succ := my_nat.succ <$> begin dedup, exact _match args.head end
-- --   | exc := tactic.fail format!"unexpected constructor: {exc}"
-- --   end
-- -- | exc := tactic.fail format!"unexpected: {exc}"
-- -- end
-- by do {tactic.interactive.json_to_expr begin exact k end >>= tactic.exact}
-- ⟩

-- run_cmd (has_to_tactic_json.to_tactic_json (my_nat.succ $ my_nat.succ $ my_nat.zero) >>= tactic.trace)

-- run_cmd (has_to_tactic_json.to_tactic_json (my_nat.succ $ my_nat.succ $ my_nat.zero) >>= (has_from_json.from_json : json → tactic my_nat) >>= tactic.trace)

@[derive [has_to_tactic_json, has_from_json]]
inductive my_tree : Type
| leaf : my_nat → my_tree
| node : my_tree → my_tree → my_nat → my_tree

meta instance : has_to_format my_tree :=
⟨λ t, match t with
| (my_tree.leaf k) := format!"(leaf {k})"
| (my_tree.node t₁ t₂ k) := by exact format!"(node {_match t₁} {_match t₂} {k})"
end
⟩

def example_tree : my_tree := my_tree.node (my_tree.leaf my_nat.zero) (my_tree.leaf my_nat.zero) my_nat.zero

-- run_cmd (has_to_tactic_json.to_tactic_json example_tree >>= tactic.trace)
-- run_cmd (has_to_tactic_json.to_tactic_json example_tree >>= (has_from_json.from_json : json → tactic my_tree) >>= tactic.trace)
end test

section instances

/-
WARNING: derived `has_to_tactic_json` and `has_from_json` instances are not guaranteed to be inverses
this can be resolved by enforcing special logic in the `has_from_json` derive handler
-/

-- meta instance : has_to_tactic_json nat :=
-- ⟨has_to_tactic_json.to_tactic_json ∘ int.of_nat⟩

-- TODO(jesse): this is insane! write the special logic.
-- attribute [derive [has_to_tactic_json, has_from_json]] nat
meta instance : has_to_tactic_json nat :=
⟨pure ∘ json.of_int ∘ int.of_nat⟩

attribute [derive has_from_json] nat -- handled by special logic in `mk_from_json`

-- run_cmd (has_to_tactic_json.to_tactic_json 3 >>=
--  λ x, tactic.trace x *> ((has_from_json.from_json : json → tactic ℕ) x >>= tactic.trace))

-- meta instance : has_from_json nat :=
-- ⟨λ msg,
--   match msg with
--   | (json.of_int (int.of_nat k)) := pure k
--   | exc := tactic.fail format!"[has_from_json_nat] unexpected: {exc}"
--   end
-- ⟩

-- meta instance : has_to_tactic_json unsigned :=
-- ⟨has_to_tactic_json.to_tactic_json ∘ unsigned.to_nat⟩

meta instance : has_from_json unsigned :=
⟨λ msg,
  match msg with
  | (json.of_int (int.of_nat k)) := pure $ unsigned.of_nat k
  | exc := tactic.fail format!"[has_from_json_unsigned] unexpected: {exc}"
  end
⟩

meta instance has_to_tactic_json_list {α} [h : has_to_tactic_json α] : has_to_tactic_json (list α) :=
⟨by mk_to_tactic_json name.anonymous⟩
universe u
-- meta instance has_from_json_list {α : Type u} [reflected α] [h : has_from_json α] : has_from_json (list α) :=
-- ⟨by mk_from_json name.anonymous⟩
-- set_option formatter.hide_full_terms false
-- run_cmd (has_to_tactic_json.to_tactic_json [1,2] >>= λ x, tactic.trace x *> (has_from_json.from_json : json → tactic (list ℕ)) x)

meta instance has_from_json_list {α} [H : has_from_json α] : has_from_json (list α) :=
⟨λ msg, do
let ⟨fn⟩ := H in
match msg with
| (json.array $ [c, json.array args]) := do
  c_nm ← has_from_json_name_aux c,
  if c_nm = `list.nil then pure [] else
  if c_nm = `list.cons then (::) <$> fn args.head <*> do x ← (args.nth 1), (by exact _match x) else
  tactic.fail format!"[has_from_json_list] unexpected {msg}"
| exc := tactic.fail "[has_from_json_list] unexpected {exc}"
end
⟩

-- run_cmd (has_to_tactic_json.to_tactic_json 2 >>= (has_from_json.from_json : json → tactic ℕ) >>= tactic.trace)

-- run_cmd (has_to_tactic_json.to_tactic_json [1,2] >>= λ x, tactic.trace x *> (has_from_json_list.from_json : json → tactic (list ℕ)) x >>= tactic.trace)

meta instance has_to_tactic_json_option {α} [has_to_tactic_json α] : has_to_tactic_json (option α) :=
⟨by mk_to_tactic_json name.anonymous⟩

meta instance has_from_json_option {α} [H : has_from_json α] : has_from_json (option α) :=
⟨λ msg, do
let ⟨fn⟩ := H in
match msg with
| (json.array $ [c, json.array args]) := do
  c_nm ← has_from_json_name_aux c,
  if c_nm = `option.none then pure none else
  if c_nm = `option.some then option.some <$> fn args.head else
  tactic.fail format!"[has_from_json_option] unexpected {msg}"
| exc := tactic.fail "[has_from_json_option] unexpected {exc}"
end
⟩

-- run_cmd (has_to_tactic_json.to_tactic_json (some [1,2]) >>= λ x, tactic.trace x *> (has_from_json.from_json : json → tactic (option $ list ℕ)) x >>= tactic.trace) -- sweet

meta instance has_to_tactic_json_prod {α β} [has_to_tactic_json α] [has_to_tactic_json β] : has_to_tactic_json (α × β) :=
⟨by mk_to_tactic_json name.anonymous⟩

meta instance has_from_json_prod {α β : Type} [H : has_from_json α] [H' : has_from_json β] : has_from_json (prod α β) :=
⟨λ msg, do
let ⟨fn₁⟩ := H in
let ⟨fn₂⟩ := H' in
match msg with
| (json.array $ [c, json.array args]) := do
  (c_nm : name) ← has_from_json_name_aux c,
  if c_nm = `prod.mk then prod.mk <$> (args.nth 0 >>= fn₁) <*> (args.nth 1 >>= fn₂) else
  tactic.fail format!"[has_from_json_prod] unexpected {msg}"
| exc := tactic.fail "[has_from_json_prod] unexpected {exc}"
end
⟩

attribute [derive [has_to_format]] binder_info
attribute [derive [has_to_tactic_json, has_from_json, has_reflect]] level
attribute [derive [has_to_tactic_json, has_from_json]] binder_info

section expr'

meta inductive expr'
| var         : nat → expr'
| sort        : level → expr'
| const       : name → list level → expr'
| mvar        (unique : name)  (pretty : name)  (type : expr') : expr'
| local_const (unique : name) (pretty : name) (bi : binder_info) (type : expr') : expr'
| app         : expr' → expr' → expr'
| lam        (var_name : name) (bi : binder_info) (var_type : expr') (body : expr') : expr'
| pi         (var_name : name) (bi : binder_info) (var_type : expr') (body : expr') : expr'
| elet       (var_name : name) (type : expr') (assignment : expr') (body : expr') : expr'

attribute [derive [has_to_tactic_json, has_from_json, has_reflect]] expr'

-- #check (by apply_instance : has_from_json expr')

attribute [derive [has_to_format]] expr'

meta def expr'.to_expr : expr' → expr
| (expr'.var k) := expr.var k
| (expr'.sort l) := (expr.sort l)
| (expr'.const n ls) := (expr.const n ls)
| (expr'.mvar un pr ty) := (expr.mvar un pr $ expr'.to_expr ty)
| (expr'.local_const un pr bi ty) := (expr.local_const un pr bi $ expr'.to_expr ty)
| (expr'.app e₁ e₂) := (expr.app (expr'.to_expr e₁) (expr'.to_expr e₂))
| (expr'.lam nm bi tp body) := (expr.lam nm bi (expr'.to_expr tp) (expr'.to_expr body))
| (expr'.pi nm bi tp body) := (expr.pi nm bi (expr'.to_expr tp) (expr'.to_expr body))
| (expr'.elet nm tp assn body) := (expr.elet nm (expr'.to_expr tp) (expr'.to_expr assn) (expr'.to_expr body))

-- meta def expr'.to_expr : expr' → tactic expr := λ x, tactic.trace "CONVERTING TO EXPR" *> (x.to_expr)

meta def expr.to_expr' : expr → tactic expr'
| (expr.var k) := pure $ expr'.var k
| (expr.sort l) := pure $ (expr'.sort l)
| (expr.const n ls) := pure $ (expr'.const n ls)
| (expr.mvar un pr ty) := (expr'.mvar un pr <$> expr.to_expr' ty)
| (expr.local_const un pr bi ty) := (expr'.local_const un pr bi <$> expr.to_expr' ty)
| (expr.app e₁ e₂) := (expr'.app <$> (expr.to_expr' e₁) <*> (expr.to_expr' e₂))
| (expr.lam nm bi tp body) := (expr'.lam nm bi <$> (expr.to_expr' tp) <*> (expr.to_expr' body))
| (expr.pi nm bi tp body) := (expr'.pi nm bi <$> (expr.to_expr' tp) <*> (expr.to_expr' body))
| (expr.elet nm tp assn body) := (expr'.elet nm <$> (expr.to_expr' tp) <*> (expr.to_expr' assn) <*> (expr.to_expr' body))
| (expr.macro md es) := tactic.fail "[expr.to_expr'] no macros allowed!"

-- @[instance, priority 9000]
-- meta def has_from_json_expr' : has_from_json expr' :=
-- ⟨λ msg, match msg with
--   | (json.object $ [("constr", c), ("args", json.array args)]) := do
--     (c_nm : name) ← has_from_json_name_aux c,
--     -- tactic.trace "NAME MSG:" *> tactic.trace nm_msg,
--     if c_nm = `expr'.const then do [nm_msg, levels_msg] ← pure args, result ← expr'.const <$> has_from_json_name_aux nm_msg <*> @has_from_json.from_json (list level) (by apply_instance : has_from_json (list level)) levels_msg, tactic.trace format!"[has_from_json_expr'] RESULT: {result}", result.to_expr.to_expr' else
--     -- tactic.fail format!"[has_from_json_expr'] unexpected: {msg}"
--     @has_from_json.from_json _ expr'.has_from_json msg
--   | exc := tactic.fail format!"[has_from_json_expr'] OH NO unexpected: {exc}"
-- end
-- ⟩

end expr'

meta instance : has_to_tactic_json expr :=
⟨λ e, e.erase_annotations.to_expr' >>= has_to_tactic_json.to_tactic_json⟩

meta instance : has_from_json expr :=
⟨λ msg, expr'.to_expr <$> (has_from_json.from_json : json → tactic expr') msg⟩

-- run_cmd (has_to_tactic_json.to_tactic_json `(2).to_expr >>= tactic.trace) -- interesting

-- #check expr
-- open tactic.interactive

-- run_cmd (has_to_tactic_json.to_tactic_json `foo.bar.baz >>= (λ x, tactic.trace x *> ((has_from_json.from_json : json → tactic name) x >>= tactic.trace)))

-- private theorem foo {p q : Prop} : p → q → ∀ r, (p ∧ q) ∨ r :=
-- λ h₁ h₂ h₃, or.inl (and.intro ‹_› ‹_›)

meta instance : has_to_tactic_json bool := ⟨λ b, pure ↑b⟩
meta instance : has_from_json bool := ⟨λ msg, match msg with
| (json.of_bool b) := pure b
| exc := tactic.fail format!"[has_from_json_bool] unexpected: {exc}"
end
⟩

end instances
-- #check tactic.set_env_core


import evaluation
import utils

namespace fairseq

section fairseq_api

meta structure CompletionRequest : Type :=
(prompt : string)
(max_tokens : int := 16)
(temperature : native.float := 1.0)
(nbest : int := 1)
(beam: int := 10)

meta def default_partial_req : CompletionRequest :=
{
  prompt := "",
  max_tokens := 6000,
  temperature := (1.0 : native.float),
  nbest := 1,
  beam := 10,
}

/-- this is responsible for validating parameters,
   e.g. ensuring floats are between 0 and 1 -/
meta instance : has_to_tactic_json CompletionRequest :=
let validate_max_tokens : int → bool := λ n, n ≤ 10000 in
let validate_float_frac : native.float → bool := λ k, 0 ≤ k ∧ k ≤ 1 in
let validate_and_return {α} (pred : α → bool) : α → tactic α := λ a, ((guard $ pred a) *> pure a <|> tactic.fail "VALIDATE_AND_RETURN FAILED") in
let validate_optional_and_return {α} (pred : α → bool) : option α → tactic (option α) := λ x, do {
  match x with
  | (some val) := some <$> validate_and_return pred val
  | none := pure none
  end
} in
let fn : CompletionRequest → tactic json := λ req, match req with
| ⟨prompt, max_tokens, temperature, nbest, beam⟩ := do
  max_tokens ← validate_and_return validate_max_tokens max_tokens,
  temperature ← validate_and_return validate_float_frac temperature,
  nbest ← validate_and_return (λ x, 0 ≤ x ∧ x ≤ (50 : int)) /- don't go overboard with the candidates -/ nbest,
  beam ← validate_and_return (λ x, 0 ≤ x ∧ x ≤ (50 : int)) /- don't go overboard with the candidates -/ beam,

  let pre_kvs : list (string × option json) := [
    ("prompt", json.of_string prompt),
    ("max_tokens", json.of_int max_tokens),
    ("temperature", json.of_float temperature),
    ("nbest", json.of_int nbest),
    ("beam", json.of_int beam)
  ],

  pure $ json.object $ pre_kvs.filter_map (λ ⟨k,mv⟩, prod.mk k <$> mv)
end
in ⟨fn⟩

meta def ENTRY_PT : string := "/Users/Yuhuai/Documents/research/scatter_transformer_fairseq/fairseq_cli/query.py"
meta def MODEL_PATH : string := "/Users/Yuhuai/Documents/research/scatter_transformer_fairseq/lean_multigoal_checkpoints/checkpoint_best.pt"
meta def DATA_PATH : string := "/Users/Yuhuai/Documents/research/scatter_transformer_fairseq/datasets/LeanMultiGoalStepSPBPE4000Bin"

meta def CompletionRequest.to_cmd (entry_pt: string) (model_path : string) (data_path : string)  : CompletionRequest → io (io.process.spawn_args)
| req@⟨prompt, max_tokens, temperature, nbest, beam⟩ := do
serialized_req ← io.run_tactic' $ has_to_tactic_json.to_tactic_json req,
pure {
  cmd := "python",
  args := [
      entry_pt
      , data_path
      , "--path"
      , model_path
      , "--sentencepiece-model"
      , data_path ++ "/model_4000_bpe.model"
      , "--json-msg"
      , json.unparse serialized_req
    ]
}

meta def serialize_ts
  (req : CompletionRequest)
  : tactic_state → tactic CompletionRequest := λ ts, do {
  ts_str ← postprocess_tactic_state ts, -- this function is responsible for replacing newlines with tabs and removing the "k goals" line
  let prompt : string :=
  ts_str,
  eval_trace format!"\n \n \n PROMPT: {prompt} \n \n \n ",
  pure {
    prompt := prompt,
    ..req}
}

meta def fairseq_api (entry_pt : string) (model_path : string) (data_path : string) : ModelAPI CompletionRequest :=

let get_predictions (response_msg : json) : option json :=
(lift_option $ do
    { (json.array choices) ← response_msg.lookup "choices" | none,
      /- `choices` is a list of {text: ..., index: ..., logprobs: ..., finish_reason: ...}-/
      texts ← choices.mmap (λ choice, choice.lookup "text"),
      pure texts
  }) in

let fn : CompletionRequest → io json := λ req, do {
  proc_cmds ← req.to_cmd entry_pt model_path data_path,
  response_raw ← io.cmd proc_cmds,
  io.put_str_ln' format!"RAW RESPONSE: {response_raw}",
  response_msg ← (lift_option $ json.parse response_raw) | io.fail' format!"[fairseq_api] JSON PARSE FAILED {response_raw}",
  (do predictions ← lift_option (get_predictions response_msg) | io.fail' format!"[fairseq_api] UNEXPECTED RESPONSE MSG: {response_msg}",
  io.put_str_ln' format!"PREDICTIONS: {predictions}",
  pure predictions) <|> pure (json.array $ [json.of_string $ format.to_string $ format!"ERROR {response_msg}"])
} in ⟨fn⟩

end fairseq_api

section fairseq_greedy_proof_search

meta def fairseq_greedy_proof_search_core
  (partial_req : fairseq.CompletionRequest)
  (entry_pt : string)
  (model_path : string)
  (data_path : string)
  (fuel := 5)
  : state_t GreedyProofSearchState tactic unit :=
greedy_proof_search_core
  (fairseq_api entry_pt model_path data_path)
    (fairseq.serialize_ts partial_req)
      (λ  msg n, run_best_beam_candidate (unwrap_lm_response $ some "[fairseq_greedy_proof_search_core]") msg n)
        (fuel)

meta def fairseq_greedy_proof_search
  (partial_req : fairseq.CompletionRequest)
  (entry_pt : string)
  (model_path : string)
  (data_path : string)
  (fuel := 5)
  (verbose := ff)
  : tactic unit :=
greedy_proof_search
  (fairseq_api entry_pt model_path data_path)
    (fairseq.serialize_ts partial_req)
      (λ  msg n, run_best_beam_candidate (unwrap_lm_response $ some "[fairseq_greedy_proof_search]") msg n)
        (fuel)
          (verbose)

end fairseq_greedy_proof_search

section test

-- example : true :=
-- begin
--   trythis "asdf",
--   sorry -- try using fairseq_greedy_proof_search here
-- end

-- example : true :=
-- begin
--   fairseq_greedy_proof_search {temperature :=1.0, nbest:=1, beam:=10, ..fairseq_api.default_partial_req}
--     fairseq_api.ENTRY_PT
--       fairseq_api.MODEL_PATH
--         fairseq_api.DATA_PATH,
-- end

-- open nat
-- example (n : ℕ) (m : ℕ) : nat.succ (n + m) = (nat.succ n + m) :=
-- begin
--   -- fairseq_greedy_proof_search {temperature :=0.7, nbest:=10, beam:=10, ..fairseq_api.default_partial_req}
--   --   fairseq_api.ENTRY_PT
--   --     fairseq_api.MODEL_PATH
--   --       fairseq_api.DATA_PATH,
-- end

-- #eval fairseq_api.CompletionRequest.to_cmd  {prompt := "true", nbest := 10, ..fairseq_api.default_partial_req} >>= io.cmd >>=io.put_str_ln
-- #eval fairseq_api.CompletionRequest.to_cmd  fairseq_api.default_partial_req >>= io.cmd >>= io.put_str_ln
-- #eval io.cmd {cmd:="echo", args:=["hello"]} >>= io.put_str_ln
-- #eval fairseq_api.serialize_ts fairseq_api.default_partial_req
end test

end fairseq


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

! This file was ported from Lean 3 source module tactic.squeeze
! leanprover-community/mathlib commit dff8393cf1d1fc152d148e13fe57452fc37d4852
! Please do not edit these lines, except to modify the commit id
! if you have ported upstream changes.
-/
import Mathbin.Control.Traversable.Basic
import Mathbin.Tactic.Simpa

/- ./././Mathport/Syntax/Translate/Tactic/Mathlib/Core.lean:38:34: unsupported: setup_tactic_parser -/
private unsafe def loc.to_string_aux : Option Name → String
  | none => "⊢"
  | some x => toString x
#align loc.to_string_aux loc.to_string_aux

/-- pretty print a `loc` -/
unsafe def loc.to_string : Loc → String
  | loc.ns [] => ""
  | loc.ns [none] => ""
  | loc.ns ls => String.join <| List.intersperse " " (" at" :: ls.map loc.to_string_aux)
  | loc.wildcard => " at *"
#align loc.to_string loc.to_string

/-- shift `pos` `n` columns to the left -/
unsafe def pos.move_left (p : Pos) (n : ℕ) : Pos
    where
  line := p.line
  column := p.column - n
#align pos.move_left pos.move_left

namespace Tactic

deriving instance DecidableEq for simp_arg_type

/-- Turn a `simp_arg_type` into a string. -/
unsafe instance simp_arg_type.has_to_string : ToString simp_arg_type :=
  ⟨fun a =>
    match a with
    | simp_arg_type.all_hyps => "*"
    | simp_arg_type.except n => "-" ++ toString n
    | simp_arg_type.expr e => toString e
    | simp_arg_type.symm_expr e => "←" ++ toString e⟩
#align tactic.simp_arg_type.has_to_string tactic.simp_arg_type.has_to_string

open List

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

/-- pretty print structure instance -/
unsafe def struct.to_tactic_format (e : pexpr) : tactic format := do
  let r ← e.get_structure_instance_info
  let fs ←
    zipWithM
        (fun n v => do
          let v ← to_expr v >>= pp
          pure <| f! "{n } := {v}")
        r.field_names r.field_values
  let ss := r.sources.map fun s => f! " .. {s}"
  let x : format := format.join <| List.intersperse ", " (fs ++ ss)
  pure f! " \{{x}}}"
#align tactic.struct.to_tactic_format tactic.struct.to_tactic_format

/-- Attribute containing a table that accumulates multiple `squeeze_simp` suggestions -/
@[user_attribute]
private unsafe def squeeze_loc_attr :
    user_attribute Unit (Option (List (Pos × String × List simp_arg_type × String)))
    where
  Name := `_squeeze_loc
  parser := fail "this attribute should not be used"
  descr := "table to accumulate multiple `squeeze_simp` suggestions"
#align tactic.squeeze_loc_attr tactic.squeeze_loc_attr

/-- dummy declaration used as target of `squeeze_loc` attribute -/
def squeezeLocAttrCarrier :=
  ()
#align tactic.squeeze_loc_attr_carrier Tactic.squeezeLocAttrCarrier

run_cmd
  squeeze_loc_attr.Set `` squeeze_loc_attr_carrier none true

/-- Format a list of arguments for use with `simp` and friends. This omits the
list entirely if it is empty.

Patch: `pp` was changed to `to_string` because it was getting rid of prefixes
that would be necessary for some disambiguations. -/
unsafe def render_simp_arg_list : List simp_arg_type → format
  | [] => ""
  | args => (· ++ ·) " " <| to_line_wrap_format <| args.map toString
#align tactic.render_simp_arg_list tactic.render_simp_arg_list

/-- Emit a suggestion to the user. If inside a `squeeze_scope` block,
the suggestions emitted through `mk_suggestion` will be aggregated so that
every tactic that makes a suggestion can consider multiple execution of the
same invocation.
If `at_pos` is true, make the suggestion at `p` instead of the current position. -/
unsafe def mk_suggestion (p : Pos) (pre post : String) (args : List simp_arg_type)
    (at_pos := false) : tactic Unit := do
  let xs ← squeeze_loc_attr.get_param `` squeeze_loc_attr_carrier
  match xs with
    | none => do
      let args := render_simp_arg_list args
      if at_pos then
          @scopeTrace _ p p fun _ => _root_.trace (s! "{pre }{args }{post}") (pure () : tactic Unit)
        else trace s! "{pre }{args }{post}"
    | some xs => do
      squeeze_loc_attr `` squeeze_loc_attr_carrier ((p, pre, args, post) :: xs) ff
#align tactic.mk_suggestion tactic.mk_suggestion

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

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

/-- `same_result proof tac` runs tactic `tac` and checks if the proof
produced by `tac` is equivalent to `proof`. -/
unsafe def same_result (pr : proof_state) (tac : tactic Unit) : tactic Bool := do
  let s ← get_proof_state_after tac
  pure <| some pr = s
#align tactic.same_result tactic.same_result

/-- Consumes the first list of `simp` arguments, accumulating required arguments
on the second one and unnecessary arguments on the third one.
-/
private unsafe def filter_simp_set_aux (tac : Bool → List simp_arg_type → tactic Unit)
    (args : List simp_arg_type) (pr : proof_state) :
    List simp_arg_type →
      List simp_arg_type → List simp_arg_type → tactic (List simp_arg_type × List simp_arg_type)
  | [], ys, ds => pure (ys, ds)
  | x :: xs, ys, ds => do
    let b ← same_result pr (tac true (args ++ xs ++ ys))
    if b then filter_simp_set_aux xs ys (ds x) else filter_simp_set_aux xs (ys x) ds
#align tactic.filter_simp_set_aux tactic.filter_simp_set_aux

initialize
  registerTraceClass.1 `squeeze.deleted

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

/-- make a `simp_arg_type` that references the name given as an argument -/
unsafe def name.to_simp_args (n : Name) : simp_arg_type :=
  simp_arg_type.expr <| @expr.local_const false n n default pexpr.mk_placeholder
#align tactic.name.to_simp_args tactic.name.to_simp_args

/-- If the `name` is (likely) to be overloaded, then prepend a `_root_` on it. The `expr` of an
overloaded name is constructed using `expr.macro`; this is how we guess whether it's overloaded. -/
unsafe def prepend_root_if_needed (n : Name) : tactic Name := do
  let x ← resolve_name' n
  return <|
      match x with
      | expr.macro _ _ => `_root_ ++ n
      | _ => n
#align tactic.prepend_root_if_needed tactic.prepend_root_if_needed

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

 * `slow`: adds all rfl-lemmas from the environment to the initial list (this is a slower but more
           accurate strategy)
 * `no_dflt`: did the user use the `only` keyword?
 * `args`:    list of `simp` arguments
 * `tac`:     how to invoke the underlying `simp` tactic
-/
unsafe def squeeze_simp_core (slow no_dflt : Bool) (args : List simp_arg_type)
    (tac : ∀ (no_dflt : Bool) (args : List simp_arg_type), tactic Unit)
    (mk_suggestion : List simp_arg_type → tactic Unit) : tactic Unit := do
  let v ← target >>= mk_meta_var
  let args ←
    if slow then do
        let simp_set ← attribute.get_instances `simp
        let simp_set ← simp_set.filterM <| has_attribute' `_refl_lemma
        let simp_set ← simp_set.mapM <| resolve_name' >=> pure ∘ simp_arg_type.expr
        pure <| args ++ simp_set
      else pure args
  let g ←
    retrieve do
        let g ← main_goal
        tac no_dflt args
        instantiate_mvars g
  let vs := g.list_constant'
  let vs ← vs.filterM is_simp_lemma
  let vs ← vs.mapM strip_prefix
  let vs ← vs.mapM prepend_root_if_needed
  with_local_goals' [v] (filter_simp_set tac args <| vs name.to_simp_args) >>= mk_suggestion
  tac no_dflt args
#align tactic.squeeze_simp_core tactic.squeeze_simp_core

namespace Interactive

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

Can be used as:

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

-/
unsafe def squeeze_scope (tac : itactic) : tactic Unit := do
  let none ← squeeze_loc_attr.get_param `` squeeze_loc_attr_carrier |
    pure ()
  squeeze_loc_attr `` squeeze_loc_attr_carrier (some []) ff
  finally tac do
      let some xs ← squeeze_loc_attr `` squeeze_loc_attr_carrier |
        fail "invalid state"
      let m := native.rb_lmap.of_list xs
      squeeze_loc_attr `` squeeze_loc_attr_carrier none ff
      m fun ⟨p, suggs⟩ => do
          let ⟨pre, _, post⟩ := suggs
          let suggs : List (List simp_arg_type) := suggs <| Prod.fst ∘ Prod.snd
          mk_suggestion p pre post (suggs List.union []) tt
          pure ()
#align tactic.interactive.squeeze_scope tactic.interactive.squeeze_scope

/- ./././Mathport/Syntax/Translate/Expr.lean:207:4: warning: unsupported notation `parser.optional -/
/- ./././Mathport/Syntax/Translate/Expr.lean:207:4: warning: unsupported notation `parser.optional -/
/- ./././Mathport/Syntax/Translate/Expr.lean:207:4: warning: unsupported notation `parser.optional -/
/-- `squeeze_simp`, `squeeze_simpa` and `squeeze_dsimp` perform the same
task with the difference that `squeeze_simp` relates to `simp` while
`squeeze_simpa` relates to `simpa` and `squeeze_dsimp` relates to
`dsimp`. The following applies to `squeeze_simp`, `squeeze_simpa` and
`squeeze_dsimp`.

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

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

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

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

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

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

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

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

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

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

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

/- ./././Mathport/Syntax/Translate/Expr.lean:207:4: warning: unsupported notation `parser.optional -/
/- ./././Mathport/Syntax/Translate/Expr.lean:207:4: warning: unsupported notation `parser.optional -/
/- ./././Mathport/Syntax/Translate/Expr.lean:207:4: warning: unsupported notation `parser.optional -/
/- ./././Mathport/Syntax/Translate/Expr.lean:207:4: warning: unsupported notation `parser.optional -/
/-- see `squeeze_simp` -/
unsafe def squeeze_simpa (key : parse cur_pos)
    (slow_and_accurate : parse (parser.optional (tk "?")))
    (use_iota_eqn : parse (parser.optional (tk "!"))) (no_dflt : parse only_flag)
    (hs : parse simp_arg_list) (attr_names : parse with_ident_list)
    (tgt : parse (parser.optional (tk "using" *> texpr)))
    (cfg : parse (parser.optional struct_inst)) : tactic Unit := do
  let (cfg', c) ← parse_config cfg
  let tgt' ←
    traverse
        (fun t => do
          let t ← to_expr t >>= pp
          pure f! " using {t}")
        tgt
  squeeze_simp_core slow_and_accurate no_dflt hs
      (fun l_no_dft l_args => simpa use_iota_eqn none l_no_dft l_args attr_names tgt cfg')
      fun args =>
      let use_iota_eqn := if use_iota_eqn then "!" else ""
      let attrs :=
        if attr_names then ""
        else String.join (List.intersperse " " (" with" :: attr_names toString))
      let tgt' := tgt' ""
      mk_suggestion (key 1) (s! "Try this: simpa{use_iota_eqn} only") (s! "{attrs }{tgt' }{c}") args
#align tactic.interactive.squeeze_simpa tactic.interactive.squeeze_simpa

/- ./././Mathport/Syntax/Translate/Expr.lean:207:4: warning: unsupported notation `parser.optional -/
/- ./././Mathport/Syntax/Translate/Expr.lean:207:4: warning: unsupported notation `parser.optional -/
/- ./././Mathport/Syntax/Translate/Expr.lean:207:4: warning: unsupported notation `parser.optional -/
/-- `squeeze_dsimp` behaves like `dsimp` (including all its arguments)
and prints a `dsimp only` invocation to skip the search through the
`simp` lemma list. See the doc string of `squeeze_simp` for examples.
 -/
unsafe def squeeze_dsimp (key : parse cur_pos)
    (slow_and_accurate : parse (parser.optional (tk "?")))
    (use_iota_eqn : parse (parser.optional (tk "!"))) (no_dflt : parse only_flag)
    (hs : parse simp_arg_list) (attr_names : parse with_ident_list) (locat : parse location)
    (cfg : parse (parser.optional struct_inst)) : tactic Unit := do
  let (cfg', c) ← parse_dsimp_config cfg
  squeeze_simp_core slow_and_accurate no_dflt hs
      (fun l_no_dft l_args => dsimp l_no_dft l_args attr_names locat cfg') fun args =>
      let use_iota_eqn := if use_iota_eqn then "!" else ""
      let attrs :=
        if attr_names then ""
        else String.join (List.intersperse " " (" with" :: attr_names toString))
      let loc := loc.to_string locat
      mk_suggestion (key 1) (s! "Try this: dsimp{use_iota_eqn} only") (s! "{attrs }{loc }{c}") args
#align tactic.interactive.squeeze_dsimp tactic.interactive.squeeze_dsimp

end Interactive

end Tactic

open Tactic.Interactive

add_tactic_doc
  { Name := "squeeze_simp / squeeze_simpa / squeeze_dsimp / squeeze_scope"
    category := DocCategory.tactic
    declNames := [`` squeeze_simp, `` squeeze_dsimp, `` squeeze_simpa, `` squeeze_scope]
    tags := ["simplification", "Try this"]
    inheritDescriptionFrom := `` squeeze_simp }


import Lean.Data.HashMap
import Lean.Data.HashSet

instance [BEq α] [Hashable α] [Repr α]: Repr (Lean.HashSet α) where
  reprPrec h n := h.toList.repr n

instance [BEq α] [Hashable α] [Repr (α × β)]: Repr (Lean.HashMap α β) where
  reprPrec h n := h.toList.repr n

namespace MWE1

inductive «Entity» where
  | aspect
    (name: String)
    («specializations»: List String := List.nil)
  | concept
    (name: String)
    («specializations»: List String := List.nil)
  deriving Repr

def «Entity».name (e: «Entity»): String :=
  match e with
  | .aspect n _ => n
  | .concept n _ => n

def «Entity».«specializations» (e: «Entity»): List String :=
  match e with
  | .aspect _ s => s
  | .concept _ s => s

class «Vocabulary»where
  «ownedStatements»: List «Entity»
  deriving Repr

def v : «Vocabulary» := {
  ownedStatements := [
    «Entity».aspect "base:Container",
    «Entity».concept "mission:Component" [ "base:Container" ]
  ]
}

inductive RDeclarationKind where
  | rAspect
  | rConcept
  deriving BEq, Repr

def «Entity».toKind (e: «Entity»): RDeclarationKind :=
  match e with
  | .aspect _ _   => .rAspect
  | .concept _ _  => .rConcept

inductive Exception where
  | error (message: String)
  deriving Repr

abbrev Names := Lean.HashSet String
abbrev Name2NamesMap := Lean.HashMap String Names

structure State where
  declarations : Lean.HashMap String RDeclarationKind := .empty
  aspectSpecializations : Name2NamesMap := .empty
  conceptSpecializations : Name2NamesMap := .empty
  deriving Repr

structure Context where
  vocabularies: List «Vocabulary» := .nil
  deriving Repr

abbrev MCore := EStateM Exception State
abbrev M     := ReaderT Context MCore

def EStateM.Result.getState (r: EStateM.Result Exception State α): State :=
  match r with
  | EStateM.Result.ok _ s => s
  | _ => {}

def State.appendSpecializations
  (d: String) (dts: List RDeclarationKind)
  (ds: List String)
  (coll: State → Name2NamesMap)
  (update: State → String → Names → State)
  : M Unit := do
  let s ← get
  match s.declarations.find? d with
  | some k =>
    if dts.contains k then
      let rds : Names := (coll s).findD d .empty
      let merged : Names := ds.foldl .insert rds
      let s' : State := update s d merged
      set s'
      pure ()
    else
      throw (Exception.error s!"Error: appendSpecializations: {repr d} is registered as a {repr k}, not one of {repr dts}.")
  | none =>
    throw (Exception.error s!"Error: appendSpecializations: there is no registered {repr dts}: {repr d} to append specializations to.")

def State.updateAspectSpecializations (s: State) (d: String) (ds: Names): State :=
  { s with aspectSpecializations := s.aspectSpecializations.insert d ds }

def State.appendAspectSpecializations (a: String) (as: List String): M Unit := do
  appendSpecializations a [ .rAspect ] as State.aspectSpecializations State.updateAspectSpecializations

def State.updateConceptSpecializations (s: State) (d: String) (ds: Names): State :=
  { s with conceptSpecializations := s.conceptSpecializations.insert d ds }

def State.appendConceptSpecializations (c: String) (cs: List String): M Unit := do
  appendSpecializations c [ .rAspect, .rConcept ] cs State.conceptSpecializations State.updateConceptSpecializations

def validateStatementDeclaration (e: «Entity»): M Unit := do
  let s ← get
  match s.declarations.find? e.name with
  | some ek =>
    throw (Exception.error s!"Error: declaration conflict: {repr e} is already registered as a {repr ek}.")
  | none =>
    let s := { s with declarations := s.declarations.insert e.name e.toKind }
    set s
    pure ()

def validateVocabularyStatementDeclarations: M Unit := do
  for v in (← read).vocabularies do
    for e in v.ownedStatements do
      validateStatementDeclaration e

def validateVocabularySpecialization (e: «Entity»): M Unit := do
  let s ← get
  match s.declarations.find? e.name with
  | some ek =>
    if ek == e.toKind then
      match ek with
      | .rAspect =>
        State.appendAspectSpecializations e.name e.specializations
      | .rConcept =>
        State.appendConceptSpecializations e.name e.specializations
    else
      throw (Exception.error s!"Error: declaration inconsistency: {repr e} is registered as a {repr ek}, not a {repr e.toKind}.")
  | none =>
    pure ()

def validateVocabularySpecializations: M Unit := do
  for v in (← read).vocabularies do
    for e in v.ownedStatements do
      validateVocabularySpecialization e

def c0 : Context := { vocabularies := [v] }

def s0 : State := {}

def s1 : State := EStateM.Result.getState (validateVocabularyStatementDeclarations |>.run c0 |>.run s0)
#eval s1
-- { declarations := [("base:Container", MWE.RDeclarationKind.rAspect),
--                    ("mission:Component", MWE.RDeclarationKind.rConcept)],
--   aspectSpecializations := [],
--   conceptSpecializations := [] }

def s2 : State := EStateM.Result.getState (validateVocabularySpecializations |>.run c0 |>.run s1)
#eval s2
-- { declarations := [("base:Container", MWE.RDeclarationKind.rAspect),
--                    ("mission:Component", MWE.RDeclarationKind.rConcept)],
--   aspectSpecializations := [("base:Container", [])],
--   conceptSpecializations := [("mission:Component", ["base:Container"])] }

-- All keys of the aspectSpecialization map must have a corresponding declaration as an rAspect
theorem AllAspectsSpecializationsKeysAreDeclared (s: State):
  ∀ a: String, s.aspectSpecializations.contains a → s.declarations.find? a == some .rAspect
:= by
    sorry

-- All values of the aspectSpecialization map must have a corresponding declaration as an rAspect
theorem AllAspectsSpecializationsValuesAreDeclared (s: State) :
  ∀ (a: String) (sup : String), (s.aspectSpecializations.findD a .empty).contains sup → s.declarations.find? sup == some .rAspect
:= by
    sorry

-- All keys of the conceptSpecializations map must have a corresponding declaration as an rConcept
theorem AllConceptSpecializationsKeysAreDeclared (s: State):
  ∀ a: String, s.conceptSpecializations.contains a → s.declarations.find? a == some .rConcept
:= by
    sorry

-- All values of the conceptSpecializations map must have a corresponding declaration as an rAspect or rConcept
theorem AllConceptSpecializationsValuesAreDeclared (s: State) :
  ∀ (a: String) (sup : String), (s.conceptSpecializations.findD a .empty).contains sup →
    s.declarations.find? sup == some .rAspect || s.declarations.find? sup == some .rConcept
:= by
    sorry
end MWE1


import Lean
import Lean.Elab
import Lean.Meta
import Lean.Parser
import Lean.PrettyPrinter
import Lean.PrettyPrinter.Formatter
import MLIR.AST
import MLIR.Dialects.BuiltinModel
import Lean.Parser
import Lean.Parser.Extra

open Lean
open Lean.Parser
open Lean.Elab
open Lean.Meta
open Lean.Parser
open Lean.Parser.ParserState
open Lean.PrettyPrinter
open Lean.PrettyPrinter.Formatter

open MLIR.AST

namespace MLIR.EDSL

-- | Custom parsers for balanced brackets
inductive Bracket
| Square -- []
| Round -- ()
| Curly -- {}
| Angle -- <>
deriving Inhabited, DecidableEq

instance : ToString Bracket where
   toString :=
    fun b =>
     match b with
     | .Square => "["
     | .Round => "("
     | .Curly => "{"
     | .Angle => "<"

-- TODO: remove <Tab> from quail
def isOpenBracket(c: Char): Option Bracket :=
match c with
| '(' => some .Round
| '[' => some .Square
| '{' => some .Curly
| '<' => some .Angle
| _ => none

def isCloseBracket(c: Char):Option Bracket :=
match c with
| ')' => some .Round
| ']' => some .Square
| '{' => some .Curly
| '<' => some .Angle
| _ => none

mutual

#check ParserState

#check Format

-- 'a -> symbol
-- `a -> antiquotation `(... ,(...))
partial def consumeCloseBracket(c: Bracket)
  (startPos: String.Pos)
  (i: String.Pos)
  (input: String)
  (brackets: List Bracket)
  (ctx: ParserContext)
  (s: ParserState): ParserState := Id.run do
    match brackets with
    | b::bs =>
      if b == c
      then
        if bs == []
        then
          let parser_fn := Lean.Parser.mkNodeToken `balanced_brackets startPos
          parser_fn ctx (s.setPos (input.next i)) -- consume the input here.
        else balancedBracketsFnAux startPos (input.next i) input bs ctx s
      else s.mkError $ "| found Opened `" ++ toString b ++ "` expected to close at `" ++ toString c ++ "`"
    | _ => s.mkError $ "| found Closed `" ++ toString c ++ "`, but have no opened brackets on stack"

partial def balancedBracketsFnAux (startPos: String.Pos)
  (i: String.Pos)
  (input: String)
  (bs: List Bracket) (ctx: ParserContext) (s: ParserState): ParserState :=
  if input.atEnd i
  then s.mkError "fonud EOF"
  else
  match input.get i with
  -- opening parens
  | '(' => balancedBracketsFnAux startPos (input.next i) input (Bracket.Round::bs) ctx s
  | '[' => balancedBracketsFnAux startPos (input.next i) input (Bracket.Square::bs) ctx s
  | '<' => balancedBracketsFnAux startPos (input.next i) input (Bracket.Angle::bs) ctx s
  | '{' => balancedBracketsFnAux startPos (input.next i) input (Bracket.Curly::bs) ctx s
  -- closing parens
  | ')' => consumeCloseBracket Bracket.Round startPos i input bs ctx s
  | ']' => consumeCloseBracket Bracket.Square startPos i input bs ctx s
  | '>' => consumeCloseBracket Bracket.Angle startPos i input bs ctx s
  | '}' => consumeCloseBracket Bracket.Curly startPos i input bs ctx s
  | c => balancedBracketsFnAux startPos (input.next i) input bs ctx s

end

-- | TODO: filter tab complete by type?
def balancedBracketsFnEntry (ctx: ParserContext) (s: ParserState): ParserState :=
  if ctx.input.get s.pos == '<'
  then balancedBracketsFnAux
   (startPos := s.pos)
   (i := s.pos)
   (input := ctx.input)
   (bs := [])
   ctx s
  else s.mkError "Expected '<'"

@[inline]
def balancedBrackets : Parser :=
   withAntiquot (mkAntiquot "balancedBrackets" `balancedBrackets) {
       fn := balancedBracketsFnEntry,
       info := mkAtomicInfo "balancedBrackets" : Parser
    }

#check balancedBrackets

-- Code stolen from test/WebServer/lean
@[combinator_formatter MLIR.EDSL.balancedBrackets]
def MLIR.EDSL.balancedBrackets.formatter : Formatter := pure ()

@[combinator_parenthesizer MLIR.EDSL.balancedBrackets]
def MLIR.EDSL.balancedBracketsParenthesizer : Parenthesizer := pure ()

macro "[balanced_brackets|" xs:balancedBrackets "]" : term => do
  match xs.raw[0] with
  | .atom _ val => return (Lean.quote val: TSyntax `str)
  | _  => Macro.throwError "expected balanced bracts to have atom"

def testBalancedBrackets : String := [balanced_brackets| < { xxasdasd } > ]
#print testBalancedBrackets

-- | positive and negative numbers, hex, octal
declare_syntax_cat mlir_int
syntax numLit: mlir_int

def IntToString (i: Int): String := i.repr

instance : Quote Int := ⟨fun n => Syntax.mkNumLit <| n.repr⟩

def quoteMDimension (d: Dimension): MacroM Syntax :=
  match d with
  | Dimension.Known n => do
    `(Dimension.Known $(quote n))
  | Dimension.Unknown => `(Dimension.Unknown)

def quoteMList (k: List (TSyntax `term)) (ty: TSyntax `term): MacroM (TSyntax `term) :=
  match k with
  | [] => `(@List.nil $ty)
  | (k::ks) => do
      let sks <- quoteMList ks ty
      `($k :: $sks)

-- AFFINE SYTAX
-- ============

declare_syntax_cat affine_expr
declare_syntax_cat affine_tuple
declare_syntax_cat affine_map

syntax ident : affine_expr
syntax "(" sepBy(affine_expr, ",") ")" : affine_tuple
syntax "affine_map<" affine_tuple "->" affine_tuple ">" : affine_map

syntax "[affine_expr|" affine_expr "]" : term
syntax "[affine_tuple|" affine_tuple "]" : term
syntax "[affine_map|" affine_map "]" : term
-- syntax "[affine_map|" affine_map "]" : term

macro_rules
| `([affine_expr| $xraw:ident ]) => do
  let xstr := xraw.getId.toString
  `(AffineExpr.Var $(Lean.quote xstr))

macro_rules
| `([affine_tuple| ( $xs,* ) ]) => do
   let initList  <- `(@List.nil MLIR.AST.AffineExpr)
   let argsList <- xs.getElems.foldrM
    (init := initList)
    (fun x xs => `([affine_expr| $x] :: $xs))
   `(AffineTuple.mk $argsList)

macro_rules
| `([affine_map| affine_map< $xs:affine_tuple -> $ys:affine_tuple >]) => do
  let xs' <- `([affine_tuple| $xs])
  let ys' <- `([affine_tuple| $ys])
  `(AffineMap.mk $xs' $ys' )

-- EDSL
-- ====

declare_syntax_cat mlir_bb
declare_syntax_cat mlir_region
declare_syntax_cat mlir_op
declare_syntax_cat mlir_op_args
declare_syntax_cat mlir_op_successor_args
declare_syntax_cat mlir_op_type
declare_syntax_cat mlir_op_operand
declare_syntax_cat mlir_ops
declare_syntax_cat mlir_type

-- syntax strLit mlir_op_args ":" mlir_op_type : mlir_op -- no region
--

-- EDSL OPERANDS
-- ==============

syntax "%" numLit : mlir_op_operand

syntax "%" ident : mlir_op_operand

syntax "[mlir_op_operand|" mlir_op_operand "]" : term
macro_rules
  | `([mlir_op_operand| $$($q)]) => return q
  | `([mlir_op_operand| % $x:ident]) => `(SSAVal.SSAVal $(Lean.quote (x.getId.toString)))
  | `([mlir_op_operand| % $n:num]) => `(SSAVal.SSAVal (IntToString $n))

def operand0 := [mlir_op_operand| %x]
#print operand0

def operand1 := [mlir_op_operand| %x]
#print operand1

def operand2 := [mlir_op_operand| %0]
#print operand2

-- EDSL OP-SUCCESSOR-ARGS
-- =================

-- successor-list       ::= `[` successor (`,` successor)* `]`
-- successor            ::= caret-id (`:` bb-arg-list)?

declare_syntax_cat mlir_op_successor_arg -- bb argument
syntax "^" ident : mlir_op_successor_arg -- bb argument with no operands
-- syntax "^" ident ":" "(" mlir_op_operand","* ")" : mlir_op_successor_arg

syntax "[mlir_op_successor_arg|" mlir_op_successor_arg "]" : term

macro_rules
  | `([mlir_op_successor_arg| ^ $x:ident  ]) =>
      `(BBName.mk $(Lean.quote (x.getId.toString)))

def succ0 :  BBName := ([mlir_op_successor_arg| ^bb])
#print succ0

-- EDSL MLIR TYPES
-- ===============

syntax "[mlir_type|" mlir_type "]" : term

-- TODO: Tuple and function types don't really exists (hardcoded Op notation)
/-
 syntax "(" mlir_type,* ")" : mlir_type
macro_rules
| `([mlir_type| ( $xs,* )]) => do
      let xs <- xs.getElems.mapM (fun x => `([mlir_type| $x]))
      let x <- quoteMList xs.toList (<- `(MLIRType _))
      `(MLIRType.tuple $x)

-- syntax "(" mlir_type ")" : mlir_type
-- syntax "(" mlir_type "," mlir_type ")" : mlir_type
-- | HACK: just switch to real parsing of lists
-- syntax "(" mlir_type "," mlir_type "," mlir_type ")" : mlir_type
syntax mlir_type "->" mlir_type : mlir_type
-/

syntax "{{" term "}}" : mlir_type
syntax "!" str : mlir_type
syntax "!" ident : mlir_type
syntax ident: mlir_type

set_option hygiene false in -- allow i to expand
macro_rules
  | `([mlir_type| $x:ident ]) => do
        let xstr := x.getId.toString
        if xstr == "index"
        then
          `(MLIRType.index)
        else if xstr.front == 'i' || xstr.front == 'f'
        then do
          let xstr' := xstr.drop 1
          match xstr'.toInt? with
          | some i =>
            let lit := Lean.Syntax.mkNumLit xstr'
            if xstr.front == 'i'
            then `(MLIRType.int .Signless $lit)
            else `(MLIRType.float $lit)
          | none =>
              Macro.throwError $ "cannot convert suffix of i/f to int: " ++ xstr
        else Macro.throwError $ "expected i<int> or f<int>, found: " ++ xstr

macro_rules
| `([mlir_type| ! $x:str ]) => `(MLIRType.undefined $x)

macro_rules
| `([mlir_type| ! $x:ident ]) => `(MLIRType.undefined $(Lean.quote x.getId.toString))

macro_rules
  | `([mlir_type| $$($q)]) => `($q)

def tyIndex : MLIRTy := [mlir_type| index]
#eval tyIndex

def tyUser : MLIRTy := [mlir_type| !"lz.int"]
#eval tyUser

def tyUserIdent : MLIRTy := [mlir_type| !shape.value]
#eval tyUserIdent

def tyi32NoGap : MLIRTy := [mlir_type| i32]
#eval tyi32NoGap
def tyf32NoGap : MLIRTy := [mlir_type| f32]
#eval tyf32NoGap

macro_rules
| `([mlir_type| {{ $t }} ]) => return t -- antiquot type

-- #print tyi32'

-- Uses dialect coercion empty → builtin
example : MLIRType builtin := [mlir_type| i32]

-- Uses dialect coercion empty → empty + builtin
example : MLIRType (Dialect.empty + builtin) := [mlir_type| i32]
-- More tricky: pushes coercion into the whole construction

declare_syntax_cat mlir_dimension

syntax "?" : mlir_dimension
syntax num : mlir_dimension

syntax "[mlir_dimension|" mlir_dimension "]" : term
macro_rules
| `([mlir_dimension| ?]) => `(Dimension.Unknown)
macro_rules
| `([mlir_dimension| $x:num ]) =>
    `(Dimension.Known $x)

def dim0 := [mlir_dimension| 30]
#print dim0

def dim1 := [mlir_dimension| ?]
#print dim1

-- | 1 x 2 x 3 x ..
declare_syntax_cat mlir_dimension_list
syntax (mlir_dimension "×")* mlir_type : mlir_dimension_list

def string_to_dimension (s: String): MacroM Dimension := do
  if s == "?"
  then return Dimension.Unknown
  else if s.isNat
  then return Dimension.Known s.toNat!
  else Macro.throwError ("unknown dimension: | " ++ s ++ "  |")

-- (MLIR.EDSL.«mlir_dimension_list_×_»
--  [
--    [(MLIR.EDSL.mlir_dimension_ (numLit "3")) "×"]
--    [(MLIR.EDSL.mlir_dimension_ (numLit "3")) "×"]]
 -- (MLIR.EDSL.mlir_type__ `i32))| )

-- | TODO: assert that the string we get is of the form x3x4x?x2...
-- that is, interleaved x and other stuff.
def parseTensorDimensionList (k: Syntax) : MacroM (TSyntax `term × TSyntax `term) := do

  let ty <- `([mlir_type|  $(⟨k.getArgs.back⟩)])
  let dimensions := (k.getArg 0)
  let dimensions <- dimensions.getArgs.toList.mapM (fun x =>
    `([mlir_dimension| $(⟨x.getArg 0⟩)]))
  let dimensions <- quoteMList dimensions (<- `(MLIR.AST.Dimension))
  -- Macro.throwError $ ("unknown dimension list:\n|" ++ (toString k.getArgs) ++ "|" ++ "\nDIMS: " ++ (toString dimensions) ++ " |\nTYPE: " ++ (toString ty)++ "")
  return (dimensions, ty)

  --       let xstr := dims.getId.toString
  --       let xparts := (xstr.splitOn "x").tail!
  --       let ty := xparts.getLast!
  --       let xparts := xparts.dropLast
  --       let xparts := [] ++ xparts -- TODO: add k into this list.
  --       -- Macro.throwError $ ("unknown dimension list: |" ++ (toString xparts) ++ "| )")

  --       let tyIdent := Lean.mkIdent ty
  --       -- let tyStx <- `([mlir_type|  $(quote tyIdent)])
  --       let tyStx <-  `([mlir_type|  i32])
  --       let dims <- xparts.mapM string_to_dimension
  --       let dimsStx <- quoteMList ([k] ++ (<- dims.mapM quoteMDimension))
  --       return (dimsStx, tyStx)
  -- -- | err => Macro.throwError $  ("unknown dimension list: |" ++ err.reprint.getD "???" ++ "| )")

-- === VECTOR TYPE ===
-- TODO: where is vector type syntax defined?
-- | TODO: fix bug that does not allow a trailing times.

-- static-dim-list ::= decimal-literal (`x` decimal-literal)*
-- | Encoding lookahead with notFollowedBy
declare_syntax_cat static_dim_list
syntax sepBy(numLit, "×", "×" notFollowedBy(mlir_type <|> "[")) : static_dim_list

syntax "[static_dim_list|" static_dim_list "]" : term
macro_rules
| `([static_dim_list| $[ $ns:num ]×* ]) => do
      quoteMList (ns.toList.map (⟨·.raw⟩)) (<- `(Nat))

-- vector-dim-list := (static-dim-list `x`)? (`[` static-dim-list `]` `x`)?
declare_syntax_cat vector_dim_list
syntax (static_dim_list "×" ("[" static_dim_list "]" "×")? )? : vector_dim_list
-- vector-element-type ::= float-type | integer-type | index-type
-- vector-type ::= `vector` `<` vector-dim-list vector-element-type `>`
syntax "vector" "<" vector_dim_list mlir_type ">"  : mlir_type

set_option hygiene false in -- allow i to expand
macro_rules
| `([mlir_type| vector < $[$fixed?:static_dim_list × $[ [ $scaled?:static_dim_list ] × ]? ]? $t:mlir_type  >]) => do
      let fixedDims <- match fixed? with
        | some s =>  `([static_dim_list| $s])
        | none => `((@List.nil Nat))
      let scaledDims <- match scaled? with
        | some (some s) => `([static_dim_list| $s])
        | _ => `((@List.nil Nat))
      `(builtin.vector $fixedDims $scaledDims [mlir_type| $t])

def staticDimList0 : List Nat := [static_dim_list| 1]
#reduce staticDimList0

def staticDimList1 : List Nat := [static_dim_list| 1 × 2]
#reduce staticDimList1

def vectorTy0 := [mlir_type| vector<i32>]
#print vectorTy0

def vectorTy1 := [mlir_type| vector<2 × i32>]
#print vectorTy1

def vectorTy2 := [mlir_type| vector<2 × 3 × [ 4 ] × i32>]
#print vectorTy2

-- | TODO: is this actually necessary?
-- syntax  "<" mlir_dimension_list  ">"  : mlir_type
-- macro_rules
-- | `([mlir_type|  < $dims:mlir_dimension_list  >]) => do
--     let (dims, ty) <- parseTensorDimensionList dims
--     `(MLIRType.vector $dims $ty)

-- | TODO: fix bug that does not allow a trailing times.

syntax "tensor" "<"  mlir_dimension_list  ">"  : mlir_type
macro_rules
| `([mlir_type| tensor < $dims:mlir_dimension_list  >]) => do
    let (dims, ty) <- parseTensorDimensionList dims
    `(builtin.tensor $dims $ty)

-- | TODO: this is a huge hack.
-- | TODO: I should be able to use the lower level parser to parse this cleanly?
syntax "tensor" "<"  "*" "×" mlir_type ">"  : mlir_type
syntax "tensor" "<*" "×" mlir_type ">"  : mlir_type
syntax "tensor" "<*×" mlir_type ">"  : mlir_type

macro_rules
| `([mlir_type| tensor < *× $ty:mlir_type >]) => do
    `(builtin.tensor_unranked [mlir_type| $ty])

macro_rules
| `([mlir_type| tensor < * × $ty:mlir_type >]) => do
    `(builtin.tensor_unranked [mlir_type| $ty])

macro_rules
| `([mlir_type| tensor <* × $ty:mlir_type >]) => do
    `(builtin.tensor_unranked [mlir_type| $ty])

macro_rules
| `([mlir_type| tensor <*×$ty:mlir_type >]) => do
    `(builtin.tensor_unranked [mlir_type| $ty])

-- Automatically inferred as MLIRType builtin
def tensorTy0 := [mlir_type| tensor<3×3×i32>]
#print tensorTy0

def tensorTy1 := [mlir_type| tensor< * × i32>]
#print tensorTy1

def tensorTy2 := [mlir_type| tensor< * × f32>]
#print tensorTy2

def tensorTy3 := [mlir_type| tensor<*× f32>]
#print tensorTy3

def tensorTy4 := [mlir_type| tensor<* × f32>]
#print tensorTy4

-- Basic coercion builtin → builtin + empty
example : MLIRType (builtin + Dialect.empty) := [mlir_type| tensor<* × f32>]

syntax "tensor1d" : mlir_type
macro_rules
| `([mlir_type| tensor1d ]) => do
    `(MLIRType.tensor1d)

def tensor1dTest : MLIRType empty := [mlir_type| tensor1d]

syntax "tensor2d" : mlir_type
macro_rules
| `([mlir_type| tensor2d ]) => do
    `(MLIRType.tensor2d)

def tensor2dTest : MLIRType empty := [mlir_type| tensor2d]

-- EDSL MLIR USER ATTRIBUTES
-- =========================

-- EDSL MLIR BASIC BLOCK OPERANDS
-- ==============================

declare_syntax_cat mlir_bb_operand
syntax mlir_op_operand ":" mlir_type : mlir_bb_operand

syntax "[mlir_bb_operand|" mlir_bb_operand "]" : term

macro_rules
| `([mlir_bb_operand| $name:mlir_op_operand : $ty:mlir_type ]) =>
     `( ([mlir_op_operand| $name], [mlir_type|$ty]) )

-- EDSL MLIR BASIC BLOCKS
-- ======================

syntax (mlir_op)* : mlir_ops

syntax "[mlir_op|" mlir_op "]" : term
syntax "[mlir_ops|" mlir_ops "]" : term

macro_rules
| `([mlir_ops| $[ $ops ]*  ]) => do
      let initList: TSyntax `term <- `(@List.nil (MLIR.AST.Op _))
      let l ← ops.foldrM (init := initList)
        fun x (xs: TSyntax `term) => `([mlir_op|$x] :: $xs)
      return l

macro_rules
  | `([mlir_ops| $$($q)]) => `(coe $q)

syntax  "{" ("^" ident ("(" sepBy(mlir_bb_operand, ",") ")")? ":")? mlir_ops "}" : mlir_region
syntax "[mlir_region|" mlir_region "]": term

macro_rules
| `([mlir_region| { ^ $name:ident ( $operands,* ) : $ops }  ]) => do
   let initList <- `(@List.nil (MLIR.AST.SSAVal × MLIR.AST.MLIRType _))
   let argsList <- operands.getElems.foldrM (init := initList) fun x xs => `([mlir_bb_operand| $x] :: $xs)
   let opsList <- `([mlir_ops| $ops])
   `(Region.mk $(Lean.quote (name.getId.toString)) $argsList $opsList)
| `([mlir_region| {  ^ $name:ident : $ops } ]) => do
   let opsList <- `([mlir_ops| $ops])
   `(Region.mk $(Lean.quote (name.getId.toString)) [] $opsList)
| `([mlir_region| { $ops:mlir_ops } ]) => do
   let opsList <- `([mlir_ops| $ops])
   `(Region.mk "entry" [] $opsList)

macro_rules
| `([mlir_region| $$($q) ]) => return q

-- TENSOR LITERAL
-- ==============

declare_syntax_cat mlir_tensor
syntax numLit : mlir_tensor
syntax scientificLit : mlir_tensor

syntax "[" sepBy(mlir_tensor, ",") "]" : mlir_tensor

syntax ident: mlir_tensor
syntax "[mlir_tensor|" mlir_tensor "]" : term

macro_rules
| `([mlir_tensor| $x:num ]) => `(TensorElem.int $x)

macro_rules
| `([mlir_tensor| $x:scientific ]) => `(TensorElem.float $(⟨x⟩))

macro_rules
| `([mlir_tensor| $x:ident ]) => do
      let xstr := x.getId.toString
      if xstr == "true"
      then `(TensorElem.bool true)
      else if xstr == "false"
      then `(TensorElem.bool false)
      else Macro.throwError ("unknown tensor value: |" ++ xstr ++ "|")

macro_rules
| `([mlir_tensor| [ $xs,* ] ]) => do
    let initList <- `([])
    let vals <- xs.getElems.foldlM (init := initList) fun xs x => `($xs ++ [[mlir_tensor| $x]])
    `(TensorElem.nested $vals)

def tensorValNum := [mlir_tensor| 42]
def tensorValFloat := [mlir_tensor| 0.000000]
def tensorValTrue := [mlir_tensor| true]
def tensorValFalse := [mlir_tensor| false]

-- MLIR ATTRIBUTE VALUE
-- ====================

-- | TODO: consider renaming this to mlir_attr
declare_syntax_cat mlir_attr_val
declare_syntax_cat mlir_attr_val_symbol
syntax "@" ident : mlir_attr_val_symbol
syntax "@" str : mlir_attr_val_symbol
syntax "#" ident : mlir_attr_val -- alias
syntax "#" strLit : mlir_attr_val -- aliass

syntax "#" ident "<" strLit ">" : mlir_attr_val -- opaqueAttr
syntax "#opaque<" ident "," strLit ">" ":" mlir_type : mlir_attr_val -- opaqueElementsAttr
syntax mlir_attr_val_symbol "::" mlir_attr_val_symbol : mlir_attr_val_symbol

declare_syntax_cat balanced_parens  -- syntax "#" ident "." ident "<" balanced_parens ">" : mlir_attr_val -- generic user attributes

syntax str: mlir_attr_val
syntax mlir_type : mlir_attr_val
syntax affine_map : mlir_attr_val
syntax mlir_attr_val_symbol : mlir_attr_val
syntax "-"? num (":" mlir_type)? : mlir_attr_val
syntax scientificLit (":" mlir_type)? : mlir_attr_val
syntax ident: mlir_attr_val

syntax "[" sepBy(mlir_attr_val, ",") "]" : mlir_attr_val
syntax "[mlir_attr_val|" mlir_attr_val "]" : term
syntax "[mlir_attr_val_symbol|" mlir_attr_val_symbol "]" : term

macro_rules
| `([mlir_attr_val| $$($x) ]) => `($x)

macro_rules
| `([mlir_attr_val|  $x:num ]) => `(AttrValue.int $x (MLIRType.int .Signless 64))
| `([mlir_attr_val| $x:num : $t:mlir_type]) => `(AttrValue.int $x [mlir_type| $t])
| `([mlir_attr_val| - $x:num ]) => `(AttrValue.int (- $x) (MLIRType.int .Signed 64))
| `([mlir_attr_val| - $x:num : $t:mlir_type]) => `(AttrValue.int (- $x) [mlir_type| $t])

macro_rules
| `([mlir_attr_val| true ]) => `(AttrValue.bool True)
| `([mlir_attr_val| false ]) => `(AttrValue.bool False)

macro_rules
| `([mlir_attr_val| # $dialect:ident < $opaqueData:str > ]) => do
  let dialect := Lean.quote dialect.getId.toString
  `(AttrValue.opaque_ $dialect $opaqueData)

macro_rules
| `([mlir_attr_val| #opaque< $dialect:ident, $opaqueData:str> : $t:mlir_type ]) => do
  let dialect := Lean.quote dialect.getId.toString
  `(AttrValue.opaqueElementsAttr $dialect $opaqueData $(⟨t⟩))

macro_rules
  | `([mlir_attr_val| $s:str]) => `(AttrValue.str $s)
  | `([mlir_attr_val| [ $xs,* ] ]) => do
        let initList <- `([])
        let vals <- xs.getElems.foldlM (init := initList) fun xs x => `($xs ++ [[mlir_attr_val| $x]])
        `(AttrValue.list $vals)
  | `([mlir_attr_val| $i:ident]) => `(AttrValue.type [mlir_type| $i:ident])
  | `([mlir_attr_val| $ty:mlir_type]) => `(AttrValue.type [mlir_type| $ty])

syntax "dense<" mlir_tensor  ">" ":" mlir_type : mlir_attr_val
macro_rules
| `([mlir_attr_val| dense< $v:mlir_tensor > : $t:mlir_type]) =>
    `(builtin.denseWithType [mlir_tensor| $v] [mlir_type| $t])

syntax "dense<" ">" ":" mlir_type: mlir_attr_val
macro_rules
| `([mlir_attr_val| dense< > : $t:mlir_type]) =>
    `(builtin.denseWithType TensorElem.empty [mlir_type| $t])

macro_rules
  | `([mlir_attr_val| $a:affine_map]) =>
      `(AttrValue.affine [affine_map| $a])

macro_rules
| `([mlir_attr_val_symbol| @ $x:str ]) =>
      `(AttrValue.symbol $x)

macro_rules
| `([mlir_attr_val_symbol| @ $x:ident ]) =>
      `(AttrValue.symbol $(Lean.quote x.getId.toString))

macro_rules
| `([mlir_attr_val_symbol| $x:mlir_attr_val_symbol :: $y:mlir_attr_val_symbol ]) =>
      `(AttrValue.nestedsymbol [mlir_attr_val_symbol| $x] [mlir_attr_val_symbol| $y])

macro_rules
| `([mlir_attr_val| $x:mlir_attr_val_symbol ]) => `([mlir_attr_val_symbol| $x])

def attrVal0Str : AttrVal := [mlir_attr_val| "foo"]
#reduce attrVal0Str

-- Uses dialect coercion: empty → builtin
example : AttrValue builtin := [mlir_attr_val| "foo"]
-- Uses dialect coercion: empty → empty + builtin
example : AttrValue (Dialect.empty + builtin) := [mlir_attr_val| "foo"]
-- Uses dialect coercion after building an AttrValue Dialect.empty

def attrVal1bTy : AttrValue builtin := [mlir_attr_val| i32]
#reduce attrVal1bTy

def attrVal2List : AttrValue builtin := [mlir_attr_val| ["foo", "foo"] ]
#reduce attrVal2List

def attrVal3AffineMap : AttrValue builtin := [mlir_attr_val| affine_map<(x, y) -> (y)>]
#reduce attrVal3AffineMap

def attrVal4Symbol : AttrValue builtin := [mlir_attr_val| @"foo" ]
#reduce attrVal4Symbol

def attrVal5int: AttrValue builtin := [mlir_attr_val| 42 ]
#reduce attrVal5int

def attrVal5bint: AttrVal := [mlir_attr_val| -42 ]
#reduce attrVal5bint

def attrVal6Symbol : AttrVal := [mlir_attr_val| @func_foo ]
#reduce attrVal6Symbol

def attrVal7NestedSymbol : AttrVal := [mlir_attr_val| @func_foo::@"func_bar" ]
#reduce attrVal7NestedSymbol

macro_rules
  | `([mlir_attr_val| # $a:str]) =>
      `(AttrValue.alias $a)

def attrVal8Alias : AttrVal := [mlir_attr_val| #"A" ]
#reduce attrVal8Alias

macro_rules
  | `([mlir_attr_val| # $a:ident]) =>
      `(AttrValue.alias $(Lean.quote a.getId.toString))

def attrVal9Alias : AttrVal := [mlir_attr_val| #a ]
#reduce attrVal9Alias

macro_rules
| `([mlir_attr_val|  $x:scientific ]) => `(AttrValue.float $(⟨x⟩) (MLIRType.float 64))
| `([mlir_attr_val| $x:scientific : $t:mlir_type]) => `(AttrValue.float $(⟨x⟩) [mlir_type| $t])

-- def attrVal10Float : AttrVal := [mlir_attr_val| 0.000000e+00  ]
def attrVal10Float :  AttrVal := [mlir_attr_val| 0.0023 ]
#print attrVal10Float

def attrVal11Escape :  AttrVal := [mlir_attr_val| $(attrVal10Float) ]
#print attrVal11Escape

-- The dense<> attribute requires the builtin dialect for the tensor type
def attrVal12DenseEmpty: AttrValue builtin := [mlir_attr_val| dense<> : tensor<0 × i64>]
#print attrVal12DenseEmpty

-- MLIR ATTRIBUTE
-- ===============

declare_syntax_cat mlir_attr_entry

syntax ident "=" mlir_attr_val : mlir_attr_entry
syntax strLit "=" mlir_attr_val : mlir_attr_entry
syntax ident : mlir_attr_entry

syntax "[mlir_attr_entry|" mlir_attr_entry "]" : term

-- | TODO: don't actually write an elaborator for the `ident` case. This forces
-- us to declare predefined identifiers in a controlled fashion.
macro_rules
  | `([mlir_attr_entry| $name:ident  = $v:mlir_attr_val]) =>
     `(AttrEntry.mk $(Lean.quote (name.getId.toString))  [mlir_attr_val| $v])
  | `([mlir_attr_entry| $name:str  = $v:mlir_attr_val]) =>
     `(AttrEntry.mk $name [mlir_attr_val| $v])

macro_rules
  | `([mlir_attr_entry| $name:ident]) =>
     `(AttrEntry.mk $(Lean.quote (name.getId.toString))  AttrValue.unit)

def attr0Str : AttrEntry builtin := [mlir_attr_entry| sym_name = "add"]
#print attr0Str

def attr2Escape : AttrEntry builtin :=
   let x : AttrVal := [mlir_attr_val| 42]
   [mlir_attr_entry| sym_name = $(x)]
#print attr0Str

def attr3Unit : AttrEntry builtin :=
   [mlir_attr_entry| sym_name]
#print attr3Unit

def attr4Negative : AttrEntry builtin :=
   [mlir_attr_entry| value = -1: i32]
#reduce attr4Negative

declare_syntax_cat mlir_attr_dict
syntax "{" sepBy(mlir_attr_entry, ",") "}" : mlir_attr_dict
syntax "[mlir_attr_dict|" mlir_attr_dict "]" : term

macro_rules
| `([mlir_attr_dict| {  $attrEntries,* } ]) => do
        let attrsList <- attrEntries.getElems.toList.mapM (fun x => `([mlir_attr_entry| $x]))
        let attrsList <- quoteMList attrsList (<- `(MLIR.AST.AttrEntry _))
        `(AttrDict.mk $attrsList)

def attrDict0 : AttrDict builtin := [mlir_attr_dict| {}]
def attrDict1 : AttrDict builtin := [mlir_attr_dict| {foo = "bar" }]
def attrDict2 : AttrDict builtin := [mlir_attr_dict| {foo = "bar", baz = "quux" }]

-- dict attribute val
syntax mlir_attr_dict : mlir_attr_val

macro_rules
| `([mlir_attr_val| $v:mlir_attr_dict]) => `(AttrValue.dict [mlir_attr_dict| $v])

def nestedAttrDict0 : AttrDict Dialect.empty := [mlir_attr_dict| {foo = {bar = "baz"} }]
#print nestedAttrDict0

-- MLIR OPS WITH REGIONS AND ATTRIBUTES AND BASIC BLOCK ARGS
-- =========================================================

--
#check sepBy1

-- Op with potential result
syntax
  (mlir_op_operand "=")?
  strLit "(" mlir_op_operand,* ")"
         ( "(" mlir_region,* ")" )?
         (mlir_attr_dict)?
  ":" "(" mlir_type,* ")" "->" "("mlir_type,*")" : mlir_op

macro_rules
  | `([mlir_op| $$($x) ]) => return x

macro_rules
  | `([mlir_op|
        $[ $resName = ]?
        $name:str
        ( $operandsNames,* )
        $[ ( $rgns,* ) ]?
        $[ $attrDict ]?
        : ( $operandsTypes,* ) -> ( $resTypes,* ) ]) => do

        -- TODO: Needs a consistency check that `resName=none ↔ resType=.unit`
        let res ← match resName with
        | none => `(@List.nil (MLIR.AST.TypedSSAVal _))
        | some name =>
           match resTypes.getElems with
           | #[] => Macro.throwError s!"expected to have return type since result '{resName}' exists"
           | #[resType] => `([([mlir_op_operand| $name], [mlir_type| $resType])])
           | tys => Macro.throwError s!"expected single return type, found multiple '{tys}'"

        -- TODO: Needs a consistency check that `operandsNames.length = operandsTypes.length`
        let operands: List (MacroM <| TSyntax `term) :=
          List.zipWith (fun x y => `(([mlir_op_operand| $x], [mlir_type| $y])))
          operandsNames.getElems.toList operandsTypes.getElems.toList
        let operands ← quoteMList (← operands.mapM id) (← `(MLIR.AST.TypedSSAVal _))
        let attrDict <- match attrDict with
                          | none => `(AttrDict.mk [])
                          | some dict => `([mlir_attr_dict| $dict])
        let rgnsList <- match rgns with
                  | none => `(@List.nil (MLIR.AST.Region _))
                  | some rgns => do
                    let rngs <- rgns.getElems.mapM (fun x => `([mlir_region| $x]))
                    quoteMList rngs.toList (<- `(MLIR.AST.Region _))

        `(Op.mk $name -- name
                $res -- results
                $operands -- operands
                $rgnsList -- regions
                $attrDict) -- attrs

-- Op with definite result
syntax mlir_op_operand "="
  strLit "(" mlir_op_operand,* ")"
         ( "(" mlir_region,* ")" )?
         (mlir_attr_dict)? ":" "(" mlir_type,* ")" "->" mlir_type : mlir_op

macro_rules
  | `([mlir_op|
        $resName:mlir_op_operand =
        $name:str
        ( $operandsNames,* )
        $[ ( $rgns,* ) ]?
        $[ $attrDict ]?
        : ( $operandsTypes,* ) -> $resType:mlir_type  ]) => do

        let res ←   `([([mlir_op_operand| $resName], [mlir_type| $resType])])
        -- TODO: Needs a consistency check that `operandsNames.length = operandsTypes.length`
        let operands: List (MacroM <| TSyntax `term) :=
          List.zipWith (fun x y => `(([mlir_op_operand| $x], [mlir_type| $y])))
          operandsNames.getElems.toList operandsTypes.getElems.toList
        let operands ← quoteMList (← operands.mapM id) (← `(MLIR.AST.TypedSSAVal _))
        let attrDict <- match attrDict with
                          | none => `(AttrDict.mk [])
                          | some dict => `([mlir_attr_dict| $dict])
        let rgnsList <- match rgns with
                  | none => `(@List.nil (MLIR.AST.Region _))
                  | some rgns => do
                    let rngs <- rgns.getElems.mapM (fun x => `([mlir_region| $x]))
                    quoteMList rngs.toList (<- `(MLIR.AST.Region _))

        `(Op.mk $name -- name
                $res -- results
                $operands -- operands
                $rgnsList -- regions
                $attrDict) -- attrs

def op1 : Op Dialect.empty :=
  [mlir_op| "foo"(%x, %y) : (i32, i32) -> (i32) ]
#print op1
def op2: Op builtin :=
  [mlir_op| %z = "foo"(%x, %y) : (i32, i32) -> (i32)]
#print op2

def bbop1 : SSAVal × MLIRTy := [mlir_bb_operand| %x : i32 ]
#print bbop1

def bb1NoArgs : Region builtin :=
  [mlir_region| {
     ^entry:
     "foo"(%x, %y) : (i32, i32) -> (i32)
      %z = "bar"(%x) : (i32) -> (i32)
      "std.return"(%x0) : (i42) -> ()
  }]
#print bb1NoArgs

def bb2SingleArg : Region builtin :=
  [mlir_region| {
     ^entry(%argp : i32):
     "foo"(%x, %y) : (i32, i32) -> (i32)
      %z = "bar"(%x) : (i32) -> (i32)
      "std.return"(%x0) : (i42) -> ()
  }]
#print bb2SingleArg

def bb3MultipleArgs : Region builtin :=
  [mlir_region| {
     ^entry(%argp : i32, %argq : i64):
     "foo"(%x, %y) : (i32, i32) -> (i32)
      %z = "bar"(%x) : (i32) -> (i32)
      "std.return"(%x0) : (i42) -> ()
  }]
#reduce bb3MultipleArgs

def rgn0 : Region Dialect.empty := ([mlir_region|  { }])
#print rgn0

def rgn1 : Region builtin :=
  [mlir_region|  {
    ^entry:
      "std.return"(%x0) : (i42) -> ()
  }]
#print rgn1

def rgn2 : Region builtin :=
  [mlir_region|  {
    ^entry:
      "std.return"(%x0) : (i42) -> ()
  }]
#print rgn2

-- | test what happens if we try to use an entry block with no explicit bb name
def rgn3 : Region builtin :=
  [mlir_region|  {
      "std.return"(%x0) : (i42) -> ()
  }]
#print rgn1

-- | test simple ops [no regions]
def opcall1 : Op Dialect.empty := [mlir_op| "foo" (%x, %y) : (i32, i32) -> (i32) ]
#print opcall1

def oprgn0 : Op Dialect.empty := [mlir_op|
 "func"() ({ ^entry: %x = "foo.add"() : () -> (i64) } ) : () -> ()
]
#reduce oprgn0

-- | note that this is a "full stack" example!
def opRgnAttr0 : Op builtin := [mlir_op|
 "module"() ({
  ^entry:
   "func"() ({
     ^bb0(%arg0:i32, %arg1:i32):
      %zero = "std.addi"(%arg0 , %arg1) : (i32, i16) -> (i64)
      "std.return"(%zero) : (i32) -> ()
    }){sym_name = "add"} : () -> ()
   "module_terminator"() : () -> ()
 }) : () -> ()
]
#print opRgnAttr0

-- | Builtins
-- =========

-- TODO: Move to `func` dialect
syntax
  "func" mlir_attr_val_symbol "(" mlir_bb_operand,* ")" ( "->" mlir_type )? "{"
    mlir_ops
  "}" : mlir_op

macro_rules
| `([mlir_op| func $name:mlir_attr_val_symbol ( $args,* ) $[ -> $ret:mlir_type ]? { $ops } ]) => do
     -- Make the arguments for the entry block
     let bbargs ← args.getElems.mapM (fun x => `([mlir_bb_operand| $x]))
     let bbargs ← quoteMList bbargs.toList (← `(MLIR.AST.SSAVal × MLIR.AST.MLIRType _))
     -- Make the entry block (the only block)
     let rgn ← `(Region.mk "entry" $bbargs [mlir_ops| $ops])

     -- Make the function signature
     let argTypes ← args.getElems.mapM (fun x => `(Prod.snd [mlir_bb_operand| $x]))
     let argTypes ← quoteMList argTypes.toList (← `(MLIR.AST.MLIRType _))
     let retType ← match ret with
       | none => `(MLIRType.tuple [])
       | some τ => `([mlir_type| $τ])
     let signature ← `(MLIRType.fn (MLIRType.tuple $argTypes) $retType)

     -- Make the entire operation
     let attrs ← `(AttrDict.mk [
        AttrEntry.mk "function_type" (AttrValue.type $signature),
        AttrEntry.mk "sym_name" [mlir_attr_val_symbol| $name]
      ])
     `(Op.mk "func" [] [] [$rgn] $attrs)

syntax "module" "{" mlir_op* "}" : mlir_op

macro_rules
| `([mlir_op| module { $ops* } ]) => do
     let initList <- `([Op.empty "module_terminator"])
     let ops <- ops.foldrM (init := initList) fun x xs => `([mlir_op| $x] :: $xs)
     let rgn <- `(Region.fromOps $ops)
     `(Op.mk "module" [] [] [$rgn] AttrDict.empty)

def mod1 : Op builtin := [mlir_op| module { }]
#print mod1

def mod2 : Op builtin := [mlir_op| module { "dummy.dummy"(): () -> () }]
#print mod2

--- MEMREF+TENSOR
--- =============
-- dimension-list ::= dimension-list-ranked | (`*` `x`)
-- dimension-list-ranked ::= (dimension `x`)*
-- dimension ::= `?` | decimal-literal
-- tensor-type ::= `tensor` `<` dimension-list tensor-memref-element-type `>`
-- tensor-memref-element-type ::= vector-element-type | vector-type | complex-type

-- https://mlir.llvm.org/docs/Dialects/Builtin/#memreftype
-- memref-type ::= ranked-memref-type | unranked-memref-type
-- ranked-memref-type ::= `memref` `<` dimension-list-ranked type
--                        (`,` layout-specification)? (`,` memory-space)? `>`
-- unranked-memref-type ::= `memref` `<*x` type (`,` memory-space)? `>`
-- stride-list ::= `[` (dimension (`,` dimension)*)? `]`
-- strided-layout ::= `offset:` dimension `,` `strides: ` stride-list
-- layout-specification ::= semi-affine-map | strided-layout | attribute-value
-- memory-space ::= attribute-value
-- | good example for paper.
declare_syntax_cat memref_type_stride_list
syntax "[" (mlir_dimension,*) "]" : memref_type_stride_list

declare_syntax_cat memref_type_strided_layout
syntax "offset:" mlir_dimension "," "strides:" memref_type_stride_list : memref_type_strided_layout

declare_syntax_cat memref_type_layout_specification
syntax memref_type_strided_layout : memref_type_layout_specification
syntax mlir_attr_val : memref_type_layout_specification
syntax "[memref_type_layout_specification|" memref_type_layout_specification "]" : term

macro_rules
| `([memref_type_layout_specification| $v:mlir_attr_val]) =>
    `(MemrefLayoutSpec.attr [mlir_attr_val| $v])
| `([memref_type_layout_specification| offset: $o:mlir_dimension , strides: [ $[ $ds:mlir_dimension ],* ]]) =>  do
    let ds <- ds.mapM (fun d => `([mlir_dimension| $d]))
    let ds <- quoteMList ds.toList (<- `(MLIR.AST.Dimension))
    `(MemrefLayoutSpec.stride [mlir_dimension| $o] $ds)

-- | ranked memref
syntax "memref" "<"  mlir_dimension_list ("," memref_type_layout_specification)? ("," mlir_attr_val)?  ">"  : mlir_type
macro_rules
| `([mlir_type| memref  < $dims:mlir_dimension_list $[, $layout ]? $[, $memspace]? >]) => do
    let (dims, ty) <- parseTensorDimensionList dims
    let memspace <- match memspace with
                    | some s => `(some [mlir_attr_val| $s])
                    | none => `(none)

    let layout <- match layout with
                  | some stx => `(some [memref_type_layout_specification| $stx])
                  | none => `(none)
    `(builtin.memref $dims $ty $layout $memspace)

def memrefTy0 := [mlir_type| memref<3×3×i32>]
#print memrefTy0

def memrefTy1 := [mlir_type| memref<i32>]
#print memrefTy1

def memrefTy2 := [mlir_type| memref<2 × 4 × i8, #map1>]
#print memrefTy2

def memrefTy3 := [mlir_type| memref<2 × 4 × i8, #map1, 1>]
#print memrefTy3

-- | unranked memref
-- unranked-memref-type ::= `memref` `<*x` type (`,` memory-space)? `>`
-- | TODO: Do I need two different parsers for these cases?
syntax "memref" "<"  "*" "×" mlir_type ("," mlir_attr_val)?  ">"  : mlir_type
syntax "memref" "<*" "×" mlir_type ("," mlir_attr_val)?  ">"  : mlir_type
macro_rules
| `([mlir_type| memref < * × $ty  $[, $memspace]? >]) => do
    let memspace <- match memspace with
                    | some s => `(some [mlir_attr_val| $s])
                    | none => `(none)
    `(builtin.memref_unranked [mlir_type| $ty] $memspace)

macro_rules
| `([mlir_type| memref <* × $ty  $[, $memspace]? >]) => do
    let memspace <- match memspace with
                    | some s => `(some [mlir_attr_val| $s])
                    | none => `(none)
    `(builtin.memref_unranked [mlir_type| $ty] $memspace)

def memrefTy4 := [mlir_type| memref<* × f32>]
#print memrefTy4

end MLIR.EDSL


import data.option.basic

/-!

# Model of the Prisma migration engine `devDiagnostic` command

The `devDiagnostic` RPC command acts as a wrapper around `diagnoseMigrationHistory`. Its
role is to interpret the diagnostic output, and translate it to a concrete
action to be performed by the CLI.

The corresponding control flow in the CLI should be:

1. Call `RPC devDiagnostic`. Check the output:
  - Error / BrokenMigration -> display the error (regular user-facing error, no
    CLI code should be needed)
  - Reset -> Prompt the user to reset with the provided reason. Call
    `RPC reset`, then proceed with 2.
  - CreateMigration -> proceed with 2.
2. Call `RPC applyMigrations`
3. If we have no migration name, prompt for it.
4. Check for the `--create-only` flag
  - If it was passed, call `RPC evaluateDataLoss`, show the warnings,
    `RPC createMigration`. Done.
  - Otherwise, call `RPC evaluateDataLoss`, potentially ask for confirmation,
    `RPC createMigration`, `RPC applyMigrations`. Generate the client. Done.

Implemented JSON-RPC API:

```typescript
interface DevDiagnosticInput {}

interface DevDiagnosticOutput {
  action: DevAction
}

type DevAction =
  { tag: "reset", reason: string }
  | { tag: "createMigration" }

```

-/

open except (error ok)

variables { α : Type }
universes u v

/-- The top-level RPC input type. -/
inductive DevInput : Type
| mk : DevInput

inductive ResetReason
| Drifted
| Unspecified

/-- The top-level RPC output type. -/
inductive DevOutput
| CreateMigration
| Reset : ResetReason → DevOutput
-- This manifests itself as a user-facing error output, it does need any special
-- handling in the CLI.
| BrokenMigration : string -> DevOutput

open DevOutput

inductive DriftDiagnostic : Type
| DriftDetected : string -> DriftDiagnostic
| MigrationFailedToApply : string -> DriftDiagnostic

inductive HistoryDiagnostic : Type
| DatabaseIsBehind
| MigrationDirectoryIsBehind
| HistoriesDiverge

structure DiagnoseMigrationHistoryOutput :=
mk ::
  ( drift : option DriftDiagnostic )
  ( history : option HistoryDiagnostic )
  ( failedMigrationNames : list string )
  ( editedMigrationNames : list string )
  ( errorInUnappliedMigrations : option string )
  ( hasMigrationsTable : bool )

def DiagnoseMigrationHistoryOutput.resetReason : DiagnoseMigrationHistoryOutput → option ResetReason :=
λ projectState,
if (
  ¬projectState.failedMigrationNames.is_nil ||
  ¬projectState.editedMigrationNames.is_nil
) then
  some ResetReason.Unspecified
else if ¬projectState.drift.is_none then
  some ResetReason.Drifted
else
  match projectState.history with
  | some HistoryDiagnostic.MigrationDirectoryIsBehind := some ResetReason.Unspecified
  | some HistoryDiagnostic.HistoriesDiverge := some ResetReason.Unspecified
  | _ := none
  end

def DiagnoseMigrationHistoryOutput.brokenMigration : DiagnoseMigrationHistoryOutput → option string :=
λ o,
match (o.drift, o.errorInUnappliedMigrations) with
| ⟨some (DriftDiagnostic.MigrationFailedToApply name), _⟩ := some name
| ⟨_, some name⟩ := some name
| _ := none
end

example : monad id := by apply_instance

/-- Machinery to define early returns. -/
def devState : Type → Type := except_t DevOutput id

instance devStateMonad : monad devState := by { unfold devState, apply_instance }
instance devStateMonadError : monad_except DevOutput devState := by { unfold devState, apply_instance }
instance devStateMonadRun : monad_run (except DevOutput) devState := by { unfold devState, apply_instance }

/-- Check that no migration (applied or unapplied) is broken. -/
def checkBrokenMigration : DiagnoseMigrationHistoryOutput → devState punit :=
λ state, match state.brokenMigration with
| some name := throw $ BrokenMigration name
| none := pure ()
end

/-- Check whether we have a ground for a reset. -/
def checkReset : DiagnoseMigrationHistoryOutput → devState punit :=
λ state, match state.resetReason with
| some reason := throw $ Reset reason
| none := pure ()
end

/-- The model implementation of `dev`. -/
def dev : DevInput → DiagnoseMigrationHistoryOutput → devState DevOutput :=
λ input projectState,
checkBrokenMigration projectState >>
  checkReset projectState >>
  pure CreateMigration

/-- Convenience wrapper around `dev` to make proof types more readable. -/
def runDev : DevInput → DiagnoseMigrationHistoryOutput → DevOutput :=
λ input diagnostics, match run (dev input diagnostics) with
| (error output) := output
| (ok output) := output
end

-- -- ---- ---- ---- ---- ---- ---- ---- ---- -
-- Proofs about `dev`'s model defined above. --
-- ---- ---- ---- ---- ---- ---- ---- ---- ----

/--
If the migrations are working and we should reset, we will always return
`Reset`. -/
theorem devReset :
  ∀ (input : DevInput) (projectState : DiagnoseMigrationHistoryOutput),
  projectState.brokenMigration = none →
  projectState.resetReason.is_some →
  ∃ r, runDev input projectState = Reset r :=
begin
  intros input projectState hBroken hReset,
  delta runDev dev checkBrokenMigration checkReset,
  obtain ⟨r, hSome⟩ : ∃ r, projectState.resetReason = some r, from option.is_some_iff_exists.mp hReset,
  existsi r,
  simp [hReset],
  rw [hSome, hBroken],
  refl
end

/--
Whenever we are not in a reset situation and no migration is broken, we will
return `CreateMigration`. -/
theorem devCreateMigration :
  ∀ (input : DevInput) (projectState : DiagnoseMigrationHistoryOutput),
  projectState.resetReason = none →
  projectState.brokenMigration = none →
  runDev input projectState = CreateMigration :=
begin
  intros input projectState hNoReset hNoBrokenMigration,
  delta runDev dev checkBrokenMigration checkReset,
  rw [hNoReset, hNoBrokenMigration],
  refl
end

/--
`dev` will always return an error before asking for a reset in case a
migration doesn't apply cleanly to the dev database. It never prematurely asks
for a reset. -/
theorem devBrokenMigration :
  ∀ (input : DevInput) (projectState : DiagnoseMigrationHistoryOutput) (brokenMigrationName : string),
  projectState.brokenMigration = some brokenMigrationName →
  runDev input projectState = BrokenMigration brokenMigrationName :=
begin
  rintros input projectState brokenMigrationName hBroken,
  unfold runDev dev checkBrokenMigration,
  simpa [hBroken]
end