/-
Copyright (c) 2020 Jason Rute. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Author(s): Jason Rute and Jesse Michael Han

REPL for interacting with Lean
-/
import data.nat.basic
import system.io
import tactic.core
import tactic

import utils
import .tactic_state
import basic.queue

section tactic_state
open interaction_monad.result
setup_tactic_parser

/- set this to tt to enable tracing
  `tactic.set_bool_option` is tactic state specific,
  and is difficult to globally set -/

meta def num_goals' : tactic_state → option ℕ :=
λ ts, match tactic.num_goals ts with | (success val _) := pure val | _ := none end

-- TODO(jesse): this is a hack. might be better to do this in python
meta def consume_with_parser {α} (p : lean.parser α) : string → io string := λ inp, do {
  io.run_tactic' $ do
    prod.snd <$> (lean.parser.run_with_input (with_input p inp) "")
}

-- TODO(jesse): performance
meta def consume_spaces : string → string
| arg@⟨[]⟩ := arg
| arg@⟨(x::xs)⟩ := if x = ' ' then consume_spaces ⟨xs⟩ else arg

-- WARNING: this is a hack
meta def remove_indents_with_split (c : char := '\t'): string → string := λ str,
let strs := str.split (= '\t') in
string.intercalate (⟨['\t']⟩ : string) (consume_spaces <$> strs)

meta def postprocess_tactic_state : tactic_state → tactic string := λ ts, do
  let main : tactic string := do {
    let ts_str := ts.to_format.to_string,
    tabbed_ts_str ← do {
      if (num_goals' ts).get_or_else 0 ≤ 1
      then pure $ ts_str.replace_char '\n' '\t'
      else tactic.unsafe_run_io $ (λ x, string.replace_char x '\n' '\t')
             <$> (consume_with_parser small_nat >=>
               consume_with_parser ident) ts_str},
    pure $ remove_indents_with_split '\t' tabbed_ts_str
  },
  main <|> (let msg := "[postprocess_tactic_state] WARNING: POSTPROCESSING FAILED" in tactic.trace msg *> tactic.fail msg)

end tactic_state

meta def add_open_namespace : name → tactic unit := λ nm, do
env ← tactic.get_env, tactic.set_env (env.execute_open nm)

meta def add_open_namespaces (nms : list name) : tactic unit :=
nms.mmap' add_open_namespace

section run_with_state'

namespace interaction_monad
open interaction_monad.result
meta def run_with_state' {σ₁ σ₂ : Type} {α : Type*} (state : σ₁) (tac : interaction_monad σ₁ α) : interaction_monad σ₂ α :=
λ s, match (tac state) with
     | (success val _) := success val s
     | (exception fn pos _) := exception fn pos s
     end
end interaction_monad
end run_with_state'

namespace tactic

open interaction_monad.result
meta def run (tac : tactic unit) : tactic (interaction_monad.result tactic_state unit) := do {
  σ ← get_state,
  match tac σ with
  | r@(success _ new_state) := interaction_monad.set_state new_state *> pure r
  | r@(exception fn pos new_state) := pure r
  end
}

-- meta instance has_format_result {α σ} [has_to_format σ] [has_to_format α] : has_to_format (interaction_monad.result σ α) := ⟨by mk_to_format `interaction_monad.result⟩ -- ayyy

meta instance has_to_format_tactic_result {α : Type*} [has_to_format α] : has_to_format (interaction_monad.result tactic_state α) :=
⟨λ r,
  match r with
  | (success val new_state) := format!"SUCCESS!\nNEW_STATE: {new_state}\nVAL: {val}"
  | (exception fn pos old_state) := do {
    let msg := (fn.get_or_else (λ _, format.of_string "n/a")) (),
    format!"EXCEPTION!\nMSG: {msg}\nPOS: {pos}\nOLD_STATE: {old_state}"
  }
  end
⟩

meta instance has_to_tactic_format_tactic_result {α : Type*} [has_to_format α] : has_to_tactic_format (interaction_monad.result tactic_state α) :=
⟨λ σ, pure $ has_to_format.to_format σ⟩

end tactic

section parse_eval_tac
setup_tactic_parser
open tactic

meta def parse_eval_tac
  (ps : parser_state)
  (tactic_string : string)
  : tactic (tactic unit × format) := do {
  let itactic_string := "{" ++ tactic_string ++ "}",
  texpr ← (reflected_value.expr ∘ prod.fst) <$>
    (interaction_monad.run_with_state' ps $ with_input parser.itactic_reflected itactic_string),
  prod.mk <$> (eval_expr (tactic unit) texpr) <*> has_to_tactic_format.to_tactic_format texpr
}

end parse_eval_tac

section frontend

open tactic lean lean.parser interactive
meta def read_eval_print_loop (ps : parser_state) : tactic unit :=
do
  trace "\nTactic state:",
  trace_state,
  let rest : tactic unit := do
  {trace "\nEnter a tactic command:",
  tactic_string <- tactic.unsafe_run_io $ io.get_line,
  (t, fmt) ← parse_eval_tac ps tactic_string,
  trace "",
  trace ("Running tactic:\n" ++ fmt.to_string),
  tactic.run t >>= eval_trace,  -- runs the tactic on the goal.  It is crashing
  read_eval_print_loop},  --- loops forever
  done <|> rest

--- like main_t, but sets the goal and environment to a user-supplied theorem in mathlib
--- note: if the environment contains all declarations in mathlib,
--- and olean files exist,
--- we don't need to supply the lean file as env.decl_olean will find it automatically.
meta def main_t_at : parser_state → tactic unit := λ ps, do {
  trace "enter declaration to prove",
  goal_nm_string ← tactic.unsafe_run_io $ io.get_line,
  ⟨nm, _⟩ ← interaction_monad.run_with_state' ps $ with_input ident goal_nm_string,
  env ← get_env,
  decl ← env.get nm,
  let g := decl.type,
  set_goal_to g,
  lean_file ← env.decl_olean nm,
  set_env_core $ environment.for_decl_of_imported_module lean_file nm,
  read_eval_print_loop ps
}

@[user_command]
meta def main_app_at
(meta_info : decl_meta_info) (_ : parse (tk "main_app_at")) : lean.parser unit :=
get_state >>= of_tactic ∘ (main_t_at)

end frontend

section main

setup_tactic_parser

open tactic
meta def main_t : parser_state → tactic unit := λ ps,
do
  trace "Enter a goal:",
  set_goal_to `(true),
  goal_string <- tactic.unsafe_run_io $ io.get_line,
  trace "GOAL STRING: " *> trace goal_string,
  (goal_pexpr, _) ← interaction_monad.run_with_state' ps $ with_input types.texpr goal_string,
  eval_trace ps.cur_pos,
  set_goal_to <e> goal_pexpr,
  (read_eval_print_loop ps) *> main_t ps  -- loops forever

@[user_command]
meta def main_app
(meta_info : decl_meta_info) (_ : parse (tk "main_app")) : lean.parser unit :=
(get_state >>= of_tactic ∘ main_t)

end main

section parse_tac

setup_tactic_parser

open tactic

/-- Parse a reflected interactive tactic from a string.
    The result can be evaluated to a `tactic unit` by using
    `eval_expr (tactic unit)`. -/
meta def parse_itactic_reflected (tactic_string : string) : tactic expr :=
let itactic_string := "{ " ++ tactic_string ++  " }" in
lean.parser.run $ do
  get_state >>= λ ps, of_tactic $ do
    tactic.set_env ps.env,
    -- eval_trace format!"[parse_itactic_reflected] TRYING TO PARSE {itactic_string}",
    (reflected_value.expr ∘ prod.fst) <$>
      (@interaction_monad.run_with_state' parser_state _ _ ps $
         with_input parser.itactic_reflected itactic_string)

/-- Parse an interactive tactic from a string. -/
meta def parse_itactic (tactic_string : string) : tactic (tactic unit) := do
  rtac ← parse_itactic_reflected tactic_string,
  eval_expr (tactic unit) rtac

end parse_tac

section evaluation_harness

meta def run_tac_with_tactic_state
  (tac : tactic unit)
  (ts : tactic_state)
  : tactic (result _ _) := do {
  tactic.write ts,
  pure $ tac ts
}

meta structure EvaluationInput : Type :=
(decl_nm : name)
(ts_data : tactic_state_data)
(tactic_string : string)
(open_namespaces : list name)

/- (before_state, action, after_state) tuple -/
meta structure EvaluationResult : Type :=
(before_state : tactic_state)
(tactic_string : string)
(result : result tactic_state unit)

meta instance : has_to_format EvaluationResult :=
⟨λ ⟨σ, tac_str, r⟩,
  format.join [
    -- hmm. why doesn't the highlighting work?
    format.highlight "INPUT STATE:\n" format.color.pink,
    σ.to_format,
    "\n\n",
    format.highlight "ACTION:\n" format.color.pink,
    tac_str,
    "\n\n",
    format.highlight  "RESULT:\n" format.color.pink,
    (has_to_format.to_format r)
  ]
⟩

/-- Creates an empty tactic state. -/
meta def mk_tactic_state : tactic tactic_state :=
tactic.unsafe_run_io $ io.run_tactic' $ tactic.exact `(trivial) *> tactic.read

/-- creates tactic_state_data as if we were proving the declaration
 (currently only theorems are supported) with name `decl_nm`. -/
meta def get_tsd_at_decl (decl_nm : name) : tactic tactic_state_data := do {
  env ← tactic.get_env,
  decl ← env.get decl_nm,
  mk_tactic_state >>= tactic.write,
  ts ← tactic.read,
  tactic.set_goal_to decl.type,
  result ← tactic_state_data.get,
  tactic.write ts,
  pure result
}

meta def run_evaluation : EvaluationInput → tactic EvaluationResult :=
λ ⟨decl, ts_data, tactic_string, ns⟩, (tactic.unsafe_run_io ∘ io.run_tactic') $ do {
  tac ← parse_itactic tactic_string,
  env ← get_env_at_decl decl,
  tactic.set_env_core env,
  add_open_namespaces ns,
  ts ← rebuild_tactic_state ts_data *> tactic.read,
  EvaluationResult.mk ts tactic_string <$> run_tac_with_tactic_state tac ts
}

section greedy_proof_search

meta structure ModelAPI (input_format : Type := json) : Type :=
(query : input_format → io json)

/- for testing -/
meta def dummy_api {α} : ModelAPI α :=
⟨λ _, pure $ json.of_string "[DummyAPI] FAILURE"⟩

namespace tactic
open interaction_monad interaction_monad.result

/- capture but backtrack the state -/
meta def capture' {α} (t : tactic α) : tactic (tactic_result α) :=
λ s, match t s with
| (success r s') := success (success r s') s
| (exception f p s') := success (exception f p s') s
end

end tactic

meta def tactic_hash : tactic ℕ := do {
  gs ← tactic.get_goals,
  hs ← gs.mmap $ λ g, do {
    tactic.set_goal_to g,
    es ← (::) <$> tactic.target <*> tactic.local_context,
    pure $ es.foldl (λ acc e, acc + e.hash) 0},
  pure $ hs.sum
}

meta def compare_tactic_state : tactic_state → tactic_state → tactic bool := λ ts₁ ts₂, do {
  ts ← tactic.read,
  h₁ ← (tactic.write ts₁ *> tactic_hash),
  h₂ ← (tactic.write ts₂ *> tactic_hash),
  tactic.write ts,
  pure $ h₁ = h₂
}
    -- let done_handler : state_t GreedyProofSearchState tactic bool := do {
    --   state_t.lift $ do {
    --     tactic.done,
    --     tactic.result >>= (λ x, tactic.guard_sorry x <|>
    --       eval_trace format! "[greedy_proof_search_step] WARNING: result contains sorry" *> tactic.failed),
    --     tactic.result >>= (λ x, tactic.type_check x <|>
    --       eval_trace format! "[greedy_proof_search_step] WARNING: result failed typechecking" *> tactic.failed)
    --   },

meta def validate_proof (pf : expr) : tactic unit := do {
  let tac (e : expr) : tactic unit := do {
    mk_tactic_state >>= tactic.write,
    guard (bnot pf.has_meta_var),
    tactic.guard_sorry e,
    tactic.type_check e
  },
  result ← tactic.capture' (tac pf),
  match result with
  | (interaction_monad.result.success r s') := pure ()
  | (interaction_monad.result.exception f p s') := tactic.fail "[validate_proof] ERROR: VALIDATION FAILED"
  end
}

meta def tactic_state.is_done (state : tactic_state) : tactic bool := do {
  ts ← tactic.read,
  result_flag ← do {
    tactic.write state,
    (do {
       tactic.done *> pure tt
     }) <|> pure ff
  },
  tactic.write ts,
  pure result_flag
}

meta def tactic_result.is_done {α} (tr : tactic_result α) : tactic bool := do {
  match tr with
  | (interaction_monad.result.success val state) := state.is_done
  | (interaction_monad.result.exception _ _ _) := pure ff
  end
}

meta def get_tac_and_capture_result (next_candidate : string) (timeout : ℕ := 5000) : tactic (tactic_result unit) := do {
  tac ← do {
    env ← tactic.get_env,
    eval_trace format!"[get_tac_and_capture_result] PARSING TACTIC: {next_candidate}",
    tac ← parse_itactic next_candidate,
    eval_trace format!"[get_tac_and_capture_result] PARSE SUCCESSFUL",
    tactic.set_env env,
    pure tac
  },
  eval_trace format!"[get_tac_and_capture_result] TRYING TACTIC: {next_candidate}",
  result ← tactic.capture' (tactic.try_for_time timeout $ tactic.try_for 200000 tac), -- if `tac` fails, exception is captured here
  eval_trace format!"[get_tac_and_capture_result] RESULT: {result}",

  /- use tactic state hashing to fail on no-ops modulo permutation -/
  result ← match result with
    | (interaction_monad.result.success val ts') := do {
        ts ← tactic.read,
        mcond (compare_tactic_state ts ts')
        (pure $ interaction_monad.mk_exception "tactic state no-op" none ts')
        (pure result)
      }
    | exc := pure exc
  end,
  pure result
}

meta def try_get_tac_and_capture_result (tac_string : string) (timeout : ℕ := 5000) : tactic (tactic_result unit) := do {
  get_tac_and_capture_result tac_string timeout <|> do {
    let msg : format := format!"[try_get_tac_and_capture_result] parse_itactic failed on {tac_string}",
    eval_trace msg,
    interaction_monad.mk_exception msg none <$> tactic.read
  }
}

-- caller of this function is responsible for inspecting results and stopping the search
meta def run_all_beam_candidates
  (get_candidates : json → tactic (list (string × native.float)))
  (msg : json)
  (tac_timeout : ℕ := 5000)
  : tactic (list (tactic_result unit × string × native.float) × list string) := do {

  let try_candidate_state := (list (string × native.float) × (list $ option $ tactic_result unit × string × native.float)),
  let stop : option (tactic_result unit × string × native.float) → state_t try_candidate_state tactic bool :=
    λ arg, match arg with
    | some ⟨result, candidate⟩ := do {
      state_t.lift result.is_done
    }
    | none := pure ff
    end,

  /- TODO(jesse): get rid of this state_t and just use `run_async <$> ...` instead -/
  let try_candidate : state_t try_candidate_state tactic (option $ tactic_result unit × string × native.float) := do {
    state_t.lift $ eval_trace format!"[try_candidate] ENTERING",
    ts ← state_t.lift tactic.read,
    state_t.lift $ eval_trace format!"[try_candidate] READ TACTIC STATE",
    ⟨rest, _⟩ ← state_t.get,
    match rest with
    | [] := do {
      state_t.lift $ eval_trace format!"[try_candidate] END OF LOOP",
      pure $ some ⟨interaction_monad.fail "all candidates failed" ts, "FAILURE", 0.0⟩
    }
    | (next_candidate::candidates) := do  {
      state_t.modify (λ ⟨_, rs⟩, ⟨candidates, rs⟩),
      result ← monad_lift $ try_get_tac_and_capture_result next_candidate.fst tac_timeout,
      when (interaction_monad.result.is_success $ result) $
        state_t.modify $ λ ⟨candidates, rs⟩, ⟨candidates, rs ++ [some $ ⟨result, next_candidate⟩]⟩,
      state_t.lift $ eval_trace format!"[try_candidate] CAPTURED RESULT: {result}",
      pure $ some ⟨result, next_candidate⟩
    }
    end
  },

  -- let find_successful_candidates
  --   (candidates : list (string × native.float))
  --   : tactic (list (tactic_result unit × string × native.float)) := do {
  --   tasks ← candidates.mmap (λ arg, flip prod.mk arg <$> tactic.run_async (try_get_tac_and_capture_result arg.fst : tactic $ tactic_result unit)),
  --   tactic.using_new_ref ff $ λ flag, do
  --   tasks.iterM [] $ λ acc ⟨task, tac_string, score⟩, do {
  --     mcond (tactic.read_ref flag) (pure acc) $ do {
  --       let result := task.get,
  --       if (interaction_monad.result.is_success result) then do {
  --         whenM (result.is_done) $ tactic.write_ref flag tt,
  --         pure $ acc ++ [⟨result, tac_string, score⟩]
  --       } else do {
  --         pure acc
  --       }
  --     }
  --   }
  -- },

  -- this is responsible for gracefully handling "error" JSON messages and should return an empty list of candidates
  unwrapped_candidates ← get_candidates msg,
  -- eval_trace format!"[run_all_beam_candidates] UNWRAPPED CANDIDATES: {unwrapped_candidates}",
  dedup_unwrapped_candidates ← list.dedup' unwrapped_candidates,
  eval_trace format!"[run_all_beam_candidates] DEDUP_UNWRAPPED CANDIDATES: {dedup_unwrapped_candidates}",
  let candidates := list.filter (λ x, ¬ "tidy".is_prefix_of (prod.fst x)) dedup_unwrapped_candidates,

  eval_trace format!"[run_all_beam_candidates] CANDIDATES: {candidates}",
  -- old failure callback
  -- let failure_callback := do {
  -- -- do ts ← state_t.lift tactic.read, pure ⟨interaction_monad.fail "all candidates failed" ts, "FAILURE", 0.0⟩
  -- }
  --
  successful_candidates ← (prod.snd <$> prod.snd <$> state_t.run (iterate_until try_candidate stop candidates.length $ pure none) ⟨candidates, []⟩),
  -- successful_candidates ← find_successful_candidates candidates,
  eval_trace format!"[run_all_beam_candidates] EXITING TRY_CANDIDATE LOOP",
  eval_trace format!"[run_all_beam_candidates] SUCCESSFUL CANDIDATES: {successful_candidates}",
  pure ⟨successful_candidates.filter_map id, prod.fst <$> candidates⟩
}

meta def run_best_beam_candidate
  (get_candidates : json → tactic (list string))
  (msg : json)
  (tac_timeout : ℕ := 5000)
  : tactic (tactic_result unit × string × list string) := do {
  let try_candidates : state_t (list string) tactic ((tactic_result unit) × string) := do {
      ts ← state_t.lift (tactic.read),
      (next_candidate::candidates) ← state_t.get | pure (interaction_monad.fail "all candidates failed" ts, "FAILURE"),
      state_t.modify (λ _, candidates),
      flip prod.mk next_candidate <$> (state_t.lift $ try_get_tac_and_capture_result next_candidate tac_timeout)

      -- let get_tac_and_capture_result : state_t (list string) tactic ((tactic_result unit) × string) := do {
      -- tac ← state_t.lift $ do {
      --     env ← tactic.get_env,
      --     eval_trace format!"[run_best_beam_candidate] PARSING TACTIC: {next_candidate}",
      --     tac ← parse_itactic next_candidate, -- this is the only possible point of program failure
      --     eval_trace format!"[run_best_beam_candidate] PARSE SUCCESSFUL",
      --     tactic.set_env env,
      --     pure tac
      --   },
      --   state_t.lift $ do
      --     eval_trace format!"[run_best_beam_candidate] TRYING TACTIC: {next_candidate}",
      --     result ← tactic.capture' (tactic.try_for 200000 tac), -- if `tac` fails, exception is captured here
      --     eval_trace format!"[run_best_beam_candidate] RESULT: {result}",

      --     result ← match result with
      --     | (interaction_monad.result.success val ts') := do {
      --       ts ← tactic.read,
      --       -- ite ((format.to_string $ has_to_format.to_format ts) = (format.to_string $ has_to_format.to_format ts'))
      --       mcond (compare_tactic_state ts ts')
      --        (pure $ interaction_monad.mk_exception "tactic state no-op" none ts')
      --        (pure result)
      --     }
      --     | exc := pure exc
      --     end,
      --     pure ⟨result, next_candidate⟩
      -- },
      -- get_tac_and_capture_result <|> -- handle `parse_itactic` failure here
      --   let msg : format := format!"[run_best_beam_candidate.try_candidates] parse_itactic failed on {next_candidate}" in do
      --   state_t.lift (eval_trace msg),
      --   (pure ⟨interaction_monad.mk_exception
      --     msg none ts, next_candidate⟩)
  },
  let stop : tactic_result unit × string → state_t (list string) tactic bool :=
    pure ∘ interaction_monad.result.is_success ∘ prod.fst,
  candidates ← list.filter (λ x, ¬ "tidy".is_prefix_of x) <$> (get_candidates msg >>= list.dedup),
  eval_trace format!"[run_best_beam_candidate] CANDIDATES: {candidates}",
  try_candidates_result@⟨result, tac_string⟩ ← ((prod.fst <$> state_t.run (iterate_until try_candidates stop) candidates) <|>
  do {
       old_state ← tactic.read,
       pure (prod.mk (interaction_monad.mk_exception "all candidates failed" none old_state) "all failed")
  }),
  eval_trace format!"[run_best_beam_candidate] TRY_CANDIDATES_RESULT: {try_candidates_result}",
  pure ⟨result, tac_string, candidates⟩
}

-- TODO(jesse): finalize JSON format for serialize_ts/decode_response protocol
meta structure GreedyProofSearchState : Type :=
(depth : ℕ := 0)
(tactics : list string := [])
(predictions : list (list string) := [])
(states : list tactic_state := []) -- TODO(jesse): this might make the logs extremely verbose
(success : bool := ff)
(all_failed : bool := ff)
(task_id : string := "")
(global_timeout : bool := ff)
(tac_timeout : ℕ := 5)
(fuel_exhausted : bool := ff)
(decl_goal : string := "")

attribute [derive has_to_format] GreedyProofSearchState

meta instance : has_mark_global_timeout GreedyProofSearchState :=
⟨λ σ, {global_timeout := tt, ..σ}⟩

meta instance : has_register_task_id GreedyProofSearchState :=
⟨λ σ task, {task_id := task, ..σ}⟩

meta instance : has_set_tac_timeout GreedyProofSearchState :=
⟨λ σ timeout, {tac_timeout := timeout, ..σ}⟩

meta instance : has_get_tac_timeout GreedyProofSearchState :=
⟨GreedyProofSearchState.tac_timeout⟩

meta instance : has_mark_fuel_exhausted GreedyProofSearchState :=
⟨λ σ, {fuel_exhausted := tt, ..σ}⟩

meta instance : has_register_decl_goal GreedyProofSearchState :=
⟨λ σ decl_goal, {decl_goal := decl_goal, ..σ}⟩

meta def serialize_list_string : list string → json := λ xs,
  json.array $ json.of_string <$> xs

-- we supply our own instance since the default derive handler is too verbose
meta instance : has_to_tactic_json GreedyProofSearchState :=
let fn : GreedyProofSearchState → tactic json := λ σ, do {
  (serialized_states : list string) ← do {
    σ.states.mmap postprocess_tactic_state
  },
  pure $ json.object $
  [
      ("depth", json.of_int σ.depth)
    , ("tactics", serialize_list_string σ.tactics)
    , ("predictions", json.array $ serialize_list_string <$> σ.predictions)
    -- store only pretty-printed tactic states for now, with newlines replaced by tabs
    , ("states", json.array $ json.of_string <$> serialized_states)
    , ("success", json.of_bool σ.success)
    , ("all_failed", json.of_bool σ.all_failed)
    , ("task_id", json.of_string σ.task_id)
    , ("global_timeout", json.of_bool σ.global_timeout)
    , ("fuel_exhausted", json.of_bool σ.fuel_exhausted)
    , ("decl_goal", json.of_string σ.decl_goal)
  ]
} in ⟨fn⟩

namespace GreedyProofSearchState

meta def add_tac : GreedyProofSearchState → string →  GreedyProofSearchState :=
λ σ tac, {tactics := σ.tactics ++ [tac], ..σ}

meta def bump_depth : GreedyProofSearchState → GreedyProofSearchState :=
λ σ, {depth := nat.succ σ.depth, ..σ}

meta def add_prediction : GreedyProofSearchState → list string → GreedyProofSearchState :=
λ σ pred, {predictions := σ.predictions ++ [pred], ..σ}

meta def add_state : GreedyProofSearchState → tactic_state → GreedyProofSearchState :=
λ σ ts, {states := σ.states ++ [ts], ..σ}

meta def mark_all_failed : GreedyProofSearchState → GreedyProofSearchState :=
λ σ, {all_failed := tt, ..σ}

end GreedyProofSearchState

meta def greedy_proof_search_step
  {input_format : Type}
  (api : ModelAPI input_format)
  (serialize_ts : tactic_state → tactic input_format)
   -- note(jesse): this is responsible for e.g. selecting a candidate from the beam
  (decode_response : json → ℕ → tactic (tactic_result unit × string × list string))
  : state_t GreedyProofSearchState tactic (bool × GreedyProofSearchState) := do {
  state_t.lift $ eval_trace "[greedy_proof_search_step] ENTERING",
  let handle_response (response : tactic_result unit × string × list string) : state_t GreedyProofSearchState tactic unit :=
     match response with | response@⟨result, tac_string, candidates⟩ :=  do {
       if not result.is_success
       then do {
         -- let msg := format!"[greedy_proof_search_step.handle_response] UNEXPECTED
         --   interaction_monad.result.exception: {result}",
         -- eval_trace msg *> tactic.fail msg
         modify $ λ σ, σ.mark_all_failed,
         old_state ← state_t.lift $ tactic.read,
         modify $ λ σ, σ.add_state old_state,
         modify $ λ σ, σ.add_prediction candidates
       }
       else do {
         (interaction_monad.result.success _ ts) ← pure result,
         -- state_t.lift $ tactic.unsafe_run_io $ io.run_tactic' $ _
         modify $ λ σ, σ.bump_depth,
         modify $ λ σ, σ.add_tac tac_string,
         modify $ λ σ, σ.add_prediction candidates,
         old_state ← state_t.lift $ tactic.read,
         modify $ λ σ, σ.add_state old_state
       }
     }
     end,

  let get_response_and_resume : state_t GreedyProofSearchState tactic bool := do {
    tac_timeout_seconds ← GreedyProofSearchState.tac_timeout <$> get,
    response@⟨result, tac_string, candidates⟩ ← state_t.lift $ do {
      ts_msg ← tactic.read >>= serialize_ts,
      eval_trace "[greedy_proof_search_step] QUERYING API",
      response_msg ← tactic.unsafe_run_io $ api.query ts_msg,
      eval_trace format!"[greedy_proof_search_step] RESPONSE MSG {response_msg}",
      eval_trace format!"[greedy_proof_search_step] RUNNING DECODE RESPONSE WITH TIMEOUT {tac_timeout_seconds*1000}",
      response ← decode_response response_msg $ tac_timeout_seconds*1000,
      pure response
    },
    state_t.lift $ eval_trace "[greedy_proof_search_step] HANDLING RESPONSE",
    handle_response response,
    state_t.lift $ eval_trace "[greedy_proof_search_step] HANDLED RESPONSE",
    state_t.lift $ do {
      eval_trace format!"DECODED RESPONSE: {result}",
      when result.is_success $ tactic.resume result,
      eval_trace format!"RESUMED"
    },

    let done_handler : state_t GreedyProofSearchState tactic bool := do {
      state_t.lift $ do {
        tactic.done,
        tactic.result >>= (λ x, tactic.guard_sorry x <|>
          eval_trace format! "[greedy_proof_search_step] WARNING: result contains sorry" *> tactic.failed),
        tactic.result >>= (λ x, tactic.type_check x <|>
          eval_trace format! "[greedy_proof_search_step] WARNING: result failed typechecking" *> tactic.failed)
      },
      modify $ λ σ, {success := tt, ..σ},
      pure tt
    } <|> pure ff,

    let tactic_exception_handler : state_t GreedyProofSearchState tactic bool := do {
      pure (bnot result.is_success)
    },

    done_flag ← done_handler,
    exception_flag ← tactic_exception_handler,
    pure (done_flag || exception_flag)

  },

  done_flag ← get_response_and_resume,

  -- let done_handler : state_t GreedyProofSearchState tactic bool := do {
  --   state_t.lift done,
  --   modify $ λ σ, {success := tt, ..σ},
  --   pure tt
  -- },

  -- done_flag ← done_handler <|> pure ff,

  state_t.lift $ eval_trace format! "DONE FLAG: {done_flag}",

  -- pure ⟨done_flag⟩
  prod.mk done_flag <$> get  /- in second branch, no predictions succeeded and the proof search halts early -/
}

/- TODO(jesse): record proof search state + search statistics; wrap this in a state_t -/
/- `decode_response` should not modify the tactic state and currently
    should only return successful `tactic_result`s -/
/-
   in the full BFS case, `decode_response` should return a list of
   `tactic_result unit`s, which may be ```success` or `exception`
   these should then be logged by the `handle_response` function in
   `greedy_proof_search_step`, and the successful applications
   should be added to the search queue.
-/
meta def greedy_proof_search_core
  {input_format : Type}
  (api : ModelAPI input_format)
  (serialize_ts : tactic_state → tactic input_format)
  (decode_response : json → ℕ → tactic (tactic_result unit × string × list string))
  (fuel := 100000)
  : state_t GreedyProofSearchState tactic unit :=
  let fuel_exhausted_callback : state_t GreedyProofSearchState tactic (bool × GreedyProofSearchState) := do {
    state_t.modify has_mark_fuel_exhausted.mark_fuel_exhausted,
    prod.mk ff <$> get
  } in
  iterate_until
    (greedy_proof_search_step api serialize_ts decode_response)
      (λ x, pure $ x.fst = tt)
        fuel
          fuel_exhausted_callback *> pure ()

meta def greedy_proof_search
  {input_format : Type}
  (api : ModelAPI input_format)
  (serialize_ts : tactic_state → tactic input_format)
  (decode_response : json → ℕ → tactic (tactic_result unit × string × list string))
  (fuel := 100000)
  (verbose := ff)
  : tactic unit := do {
  ⟨_, σ⟩ ← state_t.run (greedy_proof_search_core api serialize_ts decode_response fuel) {},
  when verbose $ eval_trace σ
}

end greedy_proof_search

section bfs

meta structure BFSNode : Type :=
(state : tactic_state)
(score : ℤ := 0)
(tactics : list string := [])
(depth : ℕ := 0)

attribute [derive has_to_format] BFSNode

namespace BFSNode

meta def of_current_state (score : ℤ := 0) (tacs : list string := []) : tactic BFSNode := do {
  ts ← tactic.read,
  pure $ ⟨ts, score, tacs, 0⟩
}

end BFSNode

meta structure BFSState : Type :=
(depth : ℕ := 0)
(num_queries : ℕ := 0)
(tactics : list string := [])
(predictions : list (list string) := [])
(states : list tactic_state := []) -- TODO(jesse): this might make the logs extremely verbose
(success : bool := ff)
(all_failed : bool := ff)
(task_id : string := "")
(max_width : ℕ := 25) -- max qsize
(max_depth : ℕ := 50)
(nodes : pqueue BFSNode.score := pqueue.empty)
(api_failure_count : ℕ := 0)
(all_failed_count : ℕ := 0)
(global_timeout : bool := ff)
(tac_timeout : ℕ := 5)
(fuel_exhausted : bool := ff)
(decl_goal : string := "")

attribute [derive has_to_format] BFSState

meta instance : has_mark_global_timeout BFSState :=
⟨λ σ, {global_timeout := tt, ..σ}⟩

meta instance : has_register_task_id BFSState :=
⟨λ σ task, {task_id := task, ..σ}⟩

meta instance : has_set_tac_timeout BFSState :=
⟨λ σ timeout, {tac_timeout := timeout, ..σ}⟩

meta instance : has_mark_fuel_exhausted BFSState :=
⟨λ σ, {fuel_exhausted := tt, ..σ}⟩

meta instance : has_get_tac_timeout BFSState :=
⟨BFSState.tac_timeout⟩

meta instance : has_register_decl_goal BFSState :=
⟨λ σ decl_goal, {decl_goal := decl_goal, ..σ}⟩

-- meta instance : has_to_tactic_json GreedyProofSearchState :=
-- let fn : GreedyProofSearchState → tactic json := λ σ, do {
--   let serialize_list_string : list string → json := λ xs,
--     json.array $ json.of_string <$> xs,
--   (serialized_states : list string) ← do {
--     σ.states.mmap postprocess_tactic_state
--   },
--   pure $ json.object $
--   [
--       ("depth", json.of_int σ.depth)
--     , ("tactics", serialize_list_string σ.tactics)
--     , ("predictions", json.array $ serialize_list_string <$> σ.predictions)
--     -- store only pretty-printed tactic states for now, with newlines replaced by tabs
--     , ("states", json.array $ json.of_string <$> serialized_states)
--     , ("success", json.of_bool σ.success)
--     , ("all_failed", json.of_bool σ.all_failed)
--     , ("task_id", json.of_string σ.task_id)
--   ]
-- } in ⟨fn⟩

-- TODO(jesse): upgrade to full instance
meta instance : has_to_tactic_json BFSState :=
let fn : BFSState → tactic json := λ σ, do {
  (serialized_states : list string) ← do {
    σ.states.mmap $ λ x, tactic.with_full_names $ postprocess_tactic_state x
  },
  pure $ json.object
  [
      ("depth", json.of_int (σ.depth))
    , ("tactics", json.array $ (json.of_string <$> σ.tactics))
    -- , ("predictions", json.array $ serialize_list_string <$> σ.predictions)
    , ("states", json.array $ json.of_string <$> serialized_states)
    , ("success", json.of_bool (σ.success))
    , ("num_queries", json.of_int (σ.num_queries))
    , ("task_id", json.of_string $ σ.task_id)
    , ("api_failure_count", json.of_int $ σ.api_failure_count)
    , ("all_failed_count", json.of_int $ σ.all_failed_count)
    , ("global_timeout", json.of_bool (σ.global_timeout))
    , ("fuel_exhausted", json.of_bool σ.fuel_exhausted)
    , ("decl_goal", json.of_string σ.decl_goal)
  ]
}
in ⟨fn⟩

section BFSState
namespace BFSState

meta def register_task_id : BFSState → string → BFSState :=
 λ σ task_id, {task_id := task_id, ..σ}

meta def bump_num_queries : BFSState → BFSState :=
λ σ, {num_queries := nat.succ σ.num_queries, ..σ}

meta def push_nodes : BFSState → list BFSNode → BFSState :=
λ σ nodes, {nodes := σ.nodes.enqueue_many nodes, ..σ}

meta def push_node : BFSState → BFSNode → BFSState := λ S σ, S.push_nodes [σ]

meta def add_tac : BFSState → string → BFSState :=
λ σ tac, {tactics := σ.tactics ++ [tac], ..σ}

meta def add_tacs : BFSState → list string → BFSState :=
λ σ tacs, {tactics := σ.tactics ++ tacs, ..σ}

meta def bump_depth : BFSState → BFSState :=
-- λ ⟨d, tacs, preds⟩, ⟨nat.succ d, tacs, preds⟩
λ σ, {depth := nat.succ σ.depth, ..σ}

meta def add_prediction : BFSState → list string → BFSState :=
-- λ ⟨d, tacs, preds⟩ pred, ⟨d, tacs, preds ++ [pred]⟩
λ σ pred, {predictions := σ.predictions ++ [pred], ..σ}

meta def add_state : BFSState → tactic_state → BFSState :=
λ σ ts, {states := σ.states ++ [ts], ..σ}

meta def mark_all_failed : BFSState → BFSState :=
λ σ, {all_failed := tt, ..σ}

meta def mark_success : BFSState → BFSState :=
λ σ, {success := tt, ..σ}

meta def of_current_state (score : ℤ := 0) (max_width : ℕ := 25) (max_depth : ℕ := 50) (tac_timeout : ℕ := 5000) : tactic BFSState := do {
  init_node ← BFSNode.of_current_state score,
  pure {nodes := pqueue.empty.enqueue init_node, max_width := max_width, max_depth := max_depth, tac_timeout := tac_timeout}
}

meta def bump_api_failure_count : BFSState → BFSState :=
λ σ, {api_failure_count := σ.api_failure_count + 1, ..σ}

meta def bump_all_failed_count : BFSState → BFSState :=
λ σ, {all_failed_count := σ.all_failed_count + 1, ..σ}

end BFSState
end BFSState

meta def score_of_float : native.float → int :=
λ x, native.float.floor ((1000.0 : native.float) * x)

meta def pop_node : state_t BFSState tactic BFSNode := do {
  tss ← BFSState.nodes <$> get,
  match tss.dequeue with
  | (some (next, new_nodes)) := do {
      modify $ λ S, {nodes := new_nodes, ..S},
      pure next
    }
  | none := state_t.lift ∘ tactic.fail $ format!"[pop_node] pqueue empty"
  end
}

/-
  Each step of BFS does the following:
    - pops a node off the queue
    - expands the node, producing a list of descendant nodes
    - if any descendant node represents a completed proof search,
      extract the result and use it to fill the main goal
      - this is accomplished by setting the tactic state
-/

meta def bfs_step
  {input_format : Type}
  (api : ModelAPI input_format)
  (serialize_ts : tactic_state → tactic input_format)
   -- note(jesse): this is responsible for e.g. selecting a candidate from the beam
  (decode_response : json → ℕ → tactic (list (tactic_result unit × string × native.float) × list string))
  : state_t BFSState tactic (bool × BFSState) := do {
  σ ← get,
  state_t.lift $ eval_trace format!"[bfs_step] ENTERING, QUEUE STATE: {σ.nodes}",

  (some next_node) ← (some <$> pop_node <|> pure none) | (state_t.lift $ eval_trace format!"[bfs_step] queue empty") *> pure ⟨tt, σ⟩, -- yikes

  ts ← state_t.lift $ tactic.read,

  state_t.lift $ tactic.write next_node.state,

  let handle_response
    (response : list (tactic_result unit × string × native.float) × (list string))
    : state_t BFSState tactic unit :=
     match response with | response@⟨successes, candidates⟩ :=  do {
       modify $ λ σ, σ.add_prediction candidates,
       modify $ λ σ, σ.bump_num_queries,
       when (successes.length = 0) $ modify $ λ σ, σ.bump_all_failed_count,
       when (candidates.length = 0) $ modify $ λ σ, σ.bump_api_failure_count
       -- successes.mmap' $ λ ⟨result, tac_string, score⟩, do {
       --   pure ()
       -- }
     }
     end,

  let get_response_and_resume : state_t BFSState tactic bool := do {
    -- TODO(jesse): `successes` needs to be generalized from `string` to an arbitrary datatype `α`
    -- for the best-first case. for now, we just score everything 0 to get this working
    tac_timeout_seconds ← BFSState.tac_timeout <$> get,
    decl_nm ← BFSState.task_id <$> get,
    response@⟨successes, candidates⟩ ← state_t.lift $ do {
      ts_msg ← tactic.read >>= serialize_ts,
      eval_trace "[bfs_step] QUERYING API",
      response_msg ← tactic.unsafe_run_io $ api.query ts_msg,
      eval_trace format!"[bfs_step] RESPONSE MSG {response_msg}",
      eval_trace format!"[bfs_step] RUNNING DECODE RESPONSE WITH TAC_TIMEOUT {tac_timeout_seconds * 1000}",

      do {env ← tactic.get_env, d ← env.get decl_nm, tactic.trace format! "[run_proof_search_step] WARNING: GOT DECL {decl_nm}"} <|> do {tactic.trace format! "[run_proof_search_step] NO GOT DECL {decl_nm}"},
      decoded_response@⟨successes, _⟩ ← decode_response response_msg $ tac_timeout_seconds * 1000,
      eval_trace $ format!"[bfs_step] SUCCESSFUL CANDIDATES: {successes}",
      pure decoded_response
    },
    handle_response response,

    -- turn the successes into a list of new BFSNodes
    (new_nodes : list BFSNode) ← successes.mmap (λ ⟨tr, tac, score⟩, match tr with
      | (interaction_monad.result.success _ state) := do {
        -- assumes that higher float score is better
        let scale_factor := (-100000.0 : native.float),
        let int_score : int := native.float.floor $ (score * scale_factor) + next_node.score,
        pure $ BFSNode.mk state int_score (next_node.tactics ++ [tac]) (next_node.depth + 1)
      }
      | exc := state_t.lift (tactic.fail format!"[bfs_step] UNEXPECTED TACTIC RESULT WHEN CONVERTING TO new_nodes: {exc}")
      end
    ),

    monad_lift $ eval_trace format! "[bfs_step] NODES BEFORE SORTING: {new_nodes}",

    let new_nodes : list BFSNode := @list.merge_sort _ (λ x y : BFSNode, x.score < y.score) _ new_nodes,

    monad_lift $ eval_trace format! "[bfs_step] NODES AFTER SORTING: {new_nodes}",

    /- loop through the new nodes. if any of them are finished (no open goals),
       the proof search is done. otherwise, we add all of them to the BFSState pqueue
       and continue the search.
    -/

    -- modify (λ σ, σ.push_nodes new_nodes),

    -- for_ new_nodes $ λ ⟨state, score⟩, do {

    -- }
    let BFSM := state_t BFSState tactic,
    let push_node_state := (list BFSNode),

    /- note(jesse): this version of `push_node_state_tac` does not quite replicate
      the behavior of greedy_proof_search at
      max_width := 1, because it will loop over all the successful candidates anyways and check
      if the proof is finished regardless of the state of the queue
    -/

    -- let push_node_state_tac : state_t push_node_state BFSM (option BFSNode) := do {
    --   (next_node::xs) ← get | pure none,
    --   done_flag ← state_t.lift $ state_t.lift $ next_node.state.is_done,

    --   result ← if done_flag then pure (some next_node) else do {
    --     q ← BFSState.nodes <$> state_t.lift get,
    --     let qsize := q.size,
    --     state_t.lift $ state_t.lift $ eval_trace format!"[push_tac] QSIZE: {qsize}",
    --     state_t.lift $ state_t.lift $ eval_trace format!"[push_tac] QUEUE: {q}",
    --     size_exceeded ← (do σ ← state_t.lift get,
    --       pure $ to_bool $ (σ.max_width <= σ.nodes.size ∨ σ.max_depth < next_node.tactics.length)),

    --     -- TODO(jesse): more informative trace message
    --     -- TODO(jesse): currently we are erasing open tactic_states which exceed the depth limit.
    --     -- consider storing them in a pqueue of "frozen" nodes instead, and display them in the trace
    --     when size_exceeded $ monad_lift $ eval_trace format!"[push_tac] SIZE EXCEEDED",
    --     unless size_exceeded $ do {
    --       state_t.lift $ state_t.lift $ eval_trace format!"[push_tac] PUSHING NODE",
    --       state_t.lift $ state_t.modify (λ σ, σ.push_node next_node)
    --     },
    --     pure none
    --   },
    --   state_t.modify (λ _, xs),
    --   pure result
    -- },

    /-
      note(jesse): this version of `push_node_state_tac` checks that the queue is OK
      before proceeding with the rest of the loop
      currently, if max_width := 0 (this is an off-by-one misnomer), this
      reproduces the logic of greedy_proof_search.
    -/
    let push_node_state_tac : state_t push_node_state BFSM (option BFSNode) := do {
      limits_exceeded ← do {
        σ ← state_t.lift get,
        pure $ to_bool $ (σ.max_width < σ.nodes.size ∨ σ.max_depth < next_node.tactics.length)
      },

     if limits_exceeded then do {
       monad_lift $ eval_trace format!"[push_tac] SIZE EXCEEDED",
       pure none
     } else do {
       (next_node::xs) ← get | pure none,
       done_flag ← monad_lift $ next_node.state.is_done,
       if done_flag then pure (some next_node) else do {
         q ← BFSState.nodes <$> state_t.lift get,
         let qsize := q.size,
         -- monad_lift $ eval_trace format!"[push_tac] QSIZE: {qsize}",
         -- monad_lift $ eval_trace format!"[push_tac] QUEUE: {q}",
         -- monad_lift $ eval_trace format!"[push_tac] PUSHING NODE",
         state_t.lift $ do {
           state_t.modify (λ σ : BFSState, σ.push_node next_node),
           state_t.modify $ λ σ,
             if σ.depth < next_node.depth then {depth := next_node.depth, ..σ} else σ

           /- note(jesse, January 22 2021, 01:08 AM) temporarily disabling for faster evals -/
           -- state_t.modify $ λ σ, σ.add_state next_node.state -- TODO(jesse): store previous states per-node, not globally
         },
         state_t.modify (λ _, xs),
         pure none
       }
     }
    },

    (done_node) ← prod.fst <$>
      state_t.run
        (iterate_until
          push_node_state_tac
            (pure ∘ option.is_some)
              (new_nodes.length) (pure none)) new_nodes,

    match done_node with
    | (some node) := do {
      state_t.lift $ tactic.write node.state,
      state_t.modify (λ σ, σ.mark_success),
      state_t.modify (λ σ, σ.add_tacs node.tactics),
      pure tt
    }
    | none := pure ff
    end
  },
  done_flag ← get_response_and_resume,
  state_t.lift $ eval_trace format! "DONE FLAG: {done_flag}",

  -- if we are done, then the state has already been set to the successful one
  -- otherwise, backtrack to the original state
  unlessM (pure done_flag) $ state_t.lift (tactic.write ts),
  prod.mk done_flag <$> get
}

meta def bfs_core
  {input_format : Type}
  (api : ModelAPI input_format)
  (serialize_ts : tactic_state → tactic input_format)
  (decode_response : json → ℕ → tactic (list (tactic_result unit × string × native.float) × list string))
  (fuel := 100000)
  : state_t BFSState tactic unit :=
  let fuel_exhausted_callback : state_t BFSState tactic (bool × BFSState) := do {
    state_t.modify has_mark_fuel_exhausted.mark_fuel_exhausted,
    prod.mk ff <$> get
  } in do {
    ts₀ ← monad_lift $ tactic.read,
    [g] ← monad_lift $ tactic.get_goals,
    iterate_until
      (bfs_step api serialize_ts decode_response)
        (λ x, pure $ x.fst = tt)
          fuel
            fuel_exhausted_callback *> pure (),
    σ ← get,
    when σ.success $
    do {pf ← monad_lift $ tactic.get_assignment g >>= tactic.instantiate_mvars,
    (monad_lift (do {-- tactic.set_goals [g],
    pf ← tactic.get_assignment g >>= tactic.instantiate_mvars,
    tgt ← tactic.infer_type pf,
    ts₁ ← tactic.read,
    tactic.write ts₀, -- important to backtrack to the old state for validation
    validate_proof pf,
    tactic.write ts₁
    }) <|>
      (do monad_lift (eval_trace "[bfs_core] WARNING: VALIDATE_PROOF FAILED"),
       modify $ λ σ, {success := ff, ..σ}))
    }
  }

meta def bfs
  {input_format : Type}
  (api : ModelAPI input_format)
  (serialize_ts : tactic_state → tactic input_format)
  (decode_response : json → ℕ → tactic (list (tactic_result unit × string × native.float) × list string))
  (fuel := 100000)
  (verbose : bool := ff)
  (max_width : ℕ := 25)
  (max_depth : ℕ := 50)
  : tactic unit := do
  init_state ← BFSState.of_current_state 0 max_width max_depth,
  σ ← prod.snd <$> state_t.run (bfs_core api serialize_ts decode_response fuel) init_state,
  when verbose $ eval_trace σ

end bfs

namespace io.process

meta instance : has_to_format spawn_args :=
⟨λ args, args.cmd ++ " " ++ string.intercalate " " args.args⟩

end io.process

section run_proof_search

-- warning: do not use, buggy
meta def try_import (m : module_info.module_name) : tactic unit := do
let env := environment.from_imported_module_name m,
env.fold (pure () : tactic unit) (λ decl acc, acc *> tactic.try (tactic.add_decl decl))

meta def run_proof_search_step
  {search_state : Type}
  (search_core : state_t search_state tactic unit)
  (decl_nm : name)
  (open_ns : list name)
  (show_trace : bool)
  (global_timeout : ℕ)
  (init_state : tactic search_state)
  [has_mark_global_timeout search_state]
  [has_register_decl_goal search_state]
  : io (search_state × string) := do
  io.run_tactic' $ do {
    tsd ← get_tsd_at_decl decl_nm,
    eval_trace format!"[run_proof_search_step] GOT TSD AT DECL {decl_nm}",
    --- set the environment at decl
    env ← get_env_at_decl decl_nm,
    eval_trace format!"[run_proof_search_step] GOT ENV AT DECL {decl_nm}",
    tactic.set_env_core env,
    eval_trace format!"[run_proof_search_step] SET ENV AT DECL {decl_nm}",

    add_open_namespaces open_ns,
    eval_trace format!"[run_proof_search_step] ADDED OPEN NAMESPACES {open_ns}",

    /- note(jesse, January 14 2021, 09:52 AM): commented out because now we should always have namespaces in the .names files -/
    -- additionally open namespaces in decl_nm
    -- for_ decl_nm.components
    --   (λ nm, do env ← tactic.get_env, when (env.is_namespace nm) $
    --      (eval_trace format! "[run_proof_search_step] ADDING OPEN NAMESPACE {nm}") *>
    --        add_open_namespace nm),

    -- eval_trace "[run_proof_search_step] GOT TSD",

    -- eval_trace format!"[run_proof_search_step] SET ENV AT DECL {decl_nm}",
    rebuild_tactic_state tsd,
    eval_trace format!"[run_proof_search_step] REBUILT TACTIC STATE, ENTERING SEARCH CORE WITH TIMEOUT {global_timeout * 1000}",
    decl_goal_string ← format.to_string <$> (tactic.target >>= tactic.pp),
    tactic.read >>= λ ts, eval_trace format!"[run_proof_search_step] TACTIC STATE BEFORE SEARCH CORE: {ts}",
    env ← tactic.get_env,
    do {d ← env.get decl_nm, tactic.trace "[run_proof_search_step] WARNING: GOT DECL"} <|> do {tactic.trace "[run_proof_search_step] NO GOT DECL"},
    σ₀ ← flip has_register_decl_goal.register_decl_goal decl_goal_string <$> init_state,
    σ ← (prod.snd <$> tactic.try_for_time (global_timeout * 1000) (state_t.run (search_core) σ₀) <|> do {
      eval_trace format!"[run_proof_search_step] GLOBAL TIMEOUT REACHED, ABORTING ",
      pure (has_mark_global_timeout.mark_global_timeout σ₀)
    }),
    -- pure $ {task_id := decl_nm.to_string, ..σ}
    pure ⟨σ, decl_nm.to_string⟩
  }

-- to be called in an environment which imports all declarations in `mathlib`
-- ideally called on a shard of the test theorems
meta def run_proof_search_core
  {search_state : Type}
  [has_to_tactic_json search_state]
  (search_core : state_t search_state tactic unit)
  (decl_nms : list (name × list name))
  (process_search_state : search_state → io unit)
  (show_trace : bool)
  (global_timeout : ℕ) -- seconds
  -- (register_task_id : search_state → string → search_state)
  [has_register_task_id search_state]
  [has_mark_global_timeout search_state]
  [has_register_decl_goal search_state]
  (init_state : tactic search_state)
  : io unit := do {
    let run_proof_search_core_body : name → list name → io unit := λ decl_nm open_ns, do {
      when (decl_nm.length > 0) $ do
      when show_trace $ io.put_str_ln' format!"[run_proof_search_core] OUTER LOOP: PROCESSING DECL {decl_nm}",
      env₀ ← io.run_tactic' $ tactic.get_env,
      when show_trace $ io.put_str_ln' format!"[run_proof_search_core] GOT ENV",
      ⟨search_state, decl_nm⟩ ← (run_proof_search_step search_core decl_nm open_ns show_trace global_timeout init_state),
      let search_state := has_register_task_id.register_task_id search_state decl_nm,
      when show_trace $ io.put_str_ln' format!"[run_proof_search_core] GOT SEARCH STATE",
      io.run_tactic' $ tactic.set_env_core env₀,
      process_search_state search_state
    },

    for_ decl_nms (λ ⟨decl_nm, open_ns⟩, do
      run_proof_search_core_body decl_nm open_ns <|> io.put_str_ln' format!"[run_proof_search_core] WARNING: SKIPPING {decl_nm}"
)}

end run_proof_search

end evaluation_harness

section playground

-- example : ℤ → bool
-- | (of_int k) := if
-- set_option pp.notation false
-- #reduce (-3 : ℤ)

-- example (x : ℕ) : nat.zero.gcd x = x :=
-- begin
--   simp [nat.gcd]
-- end

end playground


/-
Copyright (c) 2018 Simon Hudon. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Simon Hudon, Patrick Massot
-/
import algebra.group.pi
import group_theory.group_action.defs

/-!
# Pi instances for multiplicative actions

This file defines instances for mul_action and related structures on Pi types.

## See also

* `group_theory.group_action.prod`
* `group_theory.group_action.sigma`
* `group_theory.group_action.sum`
-/

universes u v w
variable {I : Type u}     -- The indexing type
variable {f : I → Type v} -- The family of types already equipped with instances
variables (x y : Π i, f i) (i : I)

namespace pi

@[to_additive pi.has_vadd]
instance has_scalar {α : Type*} [Π i, has_scalar α $ f i] :
  has_scalar α (Π i : I, f i) :=
⟨λ s x, λ i, s • (x i)⟩

@[to_additive]
lemma smul_def {α : Type*} [Π i, has_scalar α $ f i] (s : α) : s • x = λ i, s • x i := rfl
@[simp, to_additive]
lemma smul_apply {α : Type*} [Π i, has_scalar α $ f i] (s : α) : (s • x) i = s • x i := rfl

@[to_additive pi.has_vadd']
instance has_scalar' {g : I → Type*} [Π i, has_scalar (f i) (g i)] :
  has_scalar (Π i, f i) (Π i : I, g i) :=
⟨λ s x, λ i, (s i) • (x i)⟩

@[simp, to_additive]
lemma smul_apply' {g : I → Type*} [∀ i, has_scalar (f i) (g i)] (s : Π i, f i) (x : Π i, g i) :
  (s • x) i = s i • x i :=
rfl
instance is_scalar_tower {α β : Type*}
  [has_scalar α β] [Π i, has_scalar β $ f i] [Π i, has_scalar α $ f i]
  [Π i, is_scalar_tower α β (f i)] : is_scalar_tower α β (Π i : I, f i) :=
⟨λ x y z, funext $ λ i, smul_assoc x y (z i)⟩

instance is_scalar_tower' {g : I → Type*} {α : Type*}
  [Π i, has_scalar α $ f i] [Π i, has_scalar (f i) (g i)] [Π i, has_scalar α $ g i]
  [Π i, is_scalar_tower α (f i) (g i)] : is_scalar_tower α (Π i : I, f i) (Π i : I, g i) :=
⟨λ x y z, funext $ λ i, smul_assoc x (y i) (z i)⟩

instance is_scalar_tower'' {g : I → Type*} {h : I → Type*}
  [Π i, has_scalar (f i) (g i)] [Π i, has_scalar (g i) (h i)] [Π i, has_scalar (f i) (h i)]
  [Π i, is_scalar_tower (f i) (g i) (h i)] : is_scalar_tower (Π i, f i) (Π i, g i) (Π i, h i) :=
⟨λ x y z, funext $ λ i, smul_assoc (x i) (y i) (z i)⟩

@[to_additive]
instance smul_comm_class {α β : Type*}
  [Π i, has_scalar α $ f i] [Π i, has_scalar β $ f i] [∀ i, smul_comm_class α β (f i)] :
  smul_comm_class α β (Π i : I, f i) :=
⟨λ x y z, funext $ λ i, smul_comm x y (z i)⟩

@[to_additive]
instance smul_comm_class' {g : I → Type*} {α : Type*}
  [Π i, has_scalar α $ g i] [Π i, has_scalar (f i) (g i)] [∀ i, smul_comm_class α (f i) (g i)] :
  smul_comm_class α (Π i : I, f i) (Π i : I, g i) :=
⟨λ x y z, funext $ λ i, smul_comm x (y i) (z i)⟩

@[to_additive]
instance smul_comm_class'' {g : I → Type*} {h : I → Type*}
  [Π i, has_scalar (g i) (h i)] [Π i, has_scalar (f i) (h i)]
  [∀ i, smul_comm_class (f i) (g i) (h i)] : smul_comm_class (Π i, f i) (Π i, g i) (Π i, h i) :=
⟨λ x y z, funext $ λ i, smul_comm (x i) (y i) (z i)⟩

instance {α : Type*} [Π i, has_scalar α $ f i] [Π i, has_scalar αᵐᵒᵖ $ f i]
  [∀ i, is_central_scalar α (f i)] : is_central_scalar α (Π i, f i) :=
⟨λ r m, funext $ λ i, op_smul_eq_smul _ _⟩

/-- If `f i` has a faithful scalar action for a given `i`, then so does `Π i, f i`. This is
not an instance as `i` cannot be inferred. -/
@[to_additive pi.has_faithful_vadd_at]
lemma has_faithful_smul_at {α : Type*}
  [Π i, has_scalar α $ f i] [Π i, nonempty (f i)] (i : I) [has_faithful_smul α (f i)] :
  has_faithful_smul α (Π i, f i) :=
⟨λ x y h, eq_of_smul_eq_smul $ λ a : f i, begin
  classical,
  have := congr_fun (h $ function.update (λ j, classical.choice (‹Π i, nonempty (f i)› j)) i a) i,
  simpa using this,
end⟩

@[to_additive pi.has_faithful_vadd]
instance has_faithful_smul {α : Type*}
  [nonempty I] [Π i, has_scalar α $ f i] [Π i, nonempty (f i)] [Π i, has_faithful_smul α (f i)] :
  has_faithful_smul α (Π i, f i) :=
let ⟨i⟩ := ‹nonempty I› in has_faithful_smul_at i

@[to_additive]
instance mul_action (α) {m : monoid α} [Π i, mul_action α $ f i] :
  @mul_action α (Π i : I, f i) m :=
{ smul := (•),
  mul_smul := λ r s f, funext $ λ i, mul_smul _ _ _,
  one_smul := λ f, funext $ λ i, one_smul α _ }

@[to_additive]
instance mul_action' {g : I → Type*} {m : Π i, monoid (f i)} [Π i, mul_action (f i) (g i)] :
  @mul_action (Π i, f i) (Π i : I, g i) (@pi.monoid I f m) :=
{ smul := (•),
  mul_smul := λ r s f, funext $ λ i, mul_smul _ _ _,
  one_smul := λ f, funext $ λ i, one_smul _ _ }

instance distrib_mul_action (α) {m : monoid α} {n : ∀ i, add_monoid $ f i}
  [∀ i, distrib_mul_action α $ f i] :
  @distrib_mul_action α (Π i : I, f i) m (@pi.add_monoid I f n) :=
{ smul_zero := λ c, funext $ λ i, smul_zero _,
  smul_add := λ c f g, funext $ λ i, smul_add _ _ _,
  ..pi.mul_action _ }

instance distrib_mul_action' {g : I → Type*} {m : Π i, monoid (f i)} {n : Π i, add_monoid $ g i}
  [Π i, distrib_mul_action (f i) (g i)] :
  @distrib_mul_action (Π i, f i) (Π i : I, g i) (@pi.monoid I f m) (@pi.add_monoid I g n) :=
{ smul_add := by { intros, ext x, apply smul_add },
  smul_zero := by { intros, ext x, apply smul_zero } }

lemma single_smul {α} [monoid α] [Π i, add_monoid $ f i]
  [Π i, distrib_mul_action α $ f i] [decidable_eq I] (i : I) (r : α) (x : f i) :
  single i (r • x) = r • single i x :=
single_op (λ i : I, ((•) r : f i → f i)) (λ j, smul_zero _) _ _

/-- A version of `pi.single_smul` for non-dependent functions. It is useful in cases Lean fails
to apply `pi.single_smul`. -/
lemma single_smul' {α β} [monoid α] [add_monoid β]
  [distrib_mul_action α β] [decidable_eq I] (i : I) (r : α) (x : β) :
  single i (r • x) = r • single i x :=
single_smul i r x

lemma single_smul₀ {g : I → Type*} [Π i, monoid_with_zero (f i)] [Π i, add_monoid (g i)]
  [Π i, distrib_mul_action (f i) (g i)] [decidable_eq I] (i : I) (r : f i) (x : g i) :
  single i (r • x) = single i r • single i x :=
single_op₂ (λ i : I, ((•) : f i → g i → g i)) (λ j, smul_zero _) _ _ _

instance mul_distrib_mul_action (α) {m : monoid α} {n : Π i, monoid $ f i}
  [Π i, mul_distrib_mul_action α $ f i] :
  @mul_distrib_mul_action α (Π i : I, f i) m (@pi.monoid I f n) :=
{ smul_one := λ c, funext $ λ i, smul_one _,
  smul_mul := λ c f g, funext $ λ i, smul_mul' _ _ _,
  ..pi.mul_action _ }

instance mul_distrib_mul_action' {g : I → Type*} {m : Π i, monoid (f i)} {n : Π i, monoid $ g i}
  [Π i, mul_distrib_mul_action (f i) (g i)] :
  @mul_distrib_mul_action (Π i, f i) (Π i : I, g i) (@pi.monoid I f m) (@pi.monoid I g n) :=
{ smul_mul := by { intros, ext x, apply smul_mul' },
  smul_one := by { intros, ext x, apply smul_one } }

end pi

namespace function

/-- Non-dependent version of `pi.has_scalar`. Lean gets confused by the dependent instance if this
is not present. -/
@[to_additive has_vadd]
instance has_scalar {ι R M : Type*} [has_scalar R M] :
  has_scalar R (ι → M) :=
pi.has_scalar

/-- Non-dependent version of `pi.smul_comm_class`. Lean gets confused by the dependent instance if
this is not present. -/
@[to_additive]
instance smul_comm_class {ι α β M : Type*}
  [has_scalar α M] [has_scalar β M] [smul_comm_class α β M] :
  smul_comm_class α β (ι → M) :=
pi.smul_comm_class

@[to_additive]
lemma update_smul {α : Type*} [Π i, has_scalar α (f i)] [decidable_eq I]
  (c : α) (f₁ : Π i, f i) (i : I) (x₁ : f i) :
  update (c • f₁) i (c • x₁) = c • update f₁ i x₁ :=
funext $ λ j, (apply_update (λ i, (•) c) f₁ i x₁ j).symm

end function

namespace set

@[to_additive]
lemma piecewise_smul {α : Type*} [Π i, has_scalar α (f i)] (s : set I) [Π i, decidable (i ∈ s)]
  (c : α) (f₁ g₁ : Π i, f i) :
  s.piecewise (c • f₁) (c • g₁) = c • s.piecewise f₁ g₁ :=
s.piecewise_op _ _ (λ _, (•) c)

end set

section extend

@[to_additive] lemma function.extend_smul {R α β γ : Type*} [has_scalar R γ]
  (r : R) (f : α → β) (g : α → γ) (e : β → γ) :
  function.extend f (r • g) (r • e) = r • function.extend f g e :=
funext $ λ _, by convert (apply_dite ((•) r) _ _ _).symm

end extend


import tactic
import .direction
import .boolset2d
import .listdec
import .sokostate
import .sokowidget

structure sokolevel :=
(avail : bset2d)
(ini : sokostate)
(goal : boxes_only)

instance : inhabited sokolevel
:= ⟨{ avail := [[tt]], ini := {boxes := [], storekeeper := (0,0)}, goal := ⟨[]⟩ }⟩

namespace sokolevel

def valid (level : sokolevel) := level.ini.valid level.avail
instance {level : sokolevel} : decidable level.valid
  := sokostate.valid.decidable

lemma default.valid : (default sokolevel).valid := dec_trivial

def move (level : sokolevel) (d : direction) : sokolevel
:= {ini := (level.ini.move level.avail d) ..level}

def solvable (level : sokolevel) :=
  exists goal_state : sokostate, goal_state ∈ level.goal ∧
    goal_state.reachable level.avail level.ini

lemma solvable.triv (level : sokolevel)
: level.ini ∈ level.goal → level.solvable
:= λ Hin, ⟨level.ini, ⟨Hin, sokostate.reachable.triv⟩⟩

lemma solvable.move (d : direction) (level : sokolevel)
: (level.move d).solvable → level.solvable
:= λ Hsol, exists.elim Hsol (λ goal Hsol,
  ⟨goal, ⟨Hsol.1, sokostate.reachable.move d Hsol.2⟩⟩ )

--   _                            _
--  (_)_ __ ___  _ __   ___  _ __| |_
--  | | '_ ` _ \| '_ \ / _ \| '__| __|
--  | | | | | | | |_) | (_) | |  | |_
--  |_|_| |_| |_| .__/ \___/|_|   \__|
--              |_|

def add_newline (s : sokolevel) : sokolevel := {
  avail := []::s.avail,
  ini := {
    boxes := []::s.ini.boxes,
    storekeeper := match s.ini.storekeeper with (x,y) := (x,y+1) end
  },
  goal := { boxes := []::s.goal.boxes },
}
def add_newsquare (av box stor sk : bool) (s : sokolevel) : sokolevel := {
  avail := list2d.add_to_line1 av s.avail,
  ini := {
    boxes := list2d.add_to_line1 box s.ini.boxes,
    storekeeper := if sk then (0, 0) else match s.ini.storekeeper with
      | (x,0) := (x+1,0)
      | xy := xy
    end
  },
  goal := { boxes := list2d.add_to_line1 stor s.goal.boxes, }
}

def from_string_aux : list char → option (sokolevel × bool)
| [] := some ( ⟨ [], ⟨[], (0,0)⟩, ⟨[]⟩⟩ , ff)
| (c::str) := match (from_string_aux str), c with
  | none, _ := none
  | (some (s, sk_set)), '\n' := some (s.add_newline, sk_set)
  | (some (s, sk_set)), ' ' := some (s.add_newsquare tt ff ff ff, sk_set)
  | (some (s, sk_set)), '#' := some (s.add_newsquare ff ff ff ff, sk_set)
  | (some (s, sk_set)), '.' := some (s.add_newsquare tt ff tt ff, sk_set)
  | (some (s, sk_set)), '$' := some (s.add_newsquare tt tt ff ff, sk_set)
  | (some (s, sk_set)), '*' := some (s.add_newsquare tt tt tt ff, sk_set)
  | (some (s, ff)), '@' := some (s.add_newsquare tt ff ff tt, tt)
  | (some (s, ff)), '+' := some (s.add_newsquare tt ff tt tt, tt)
  | (some _), _ := none
  end

def from_string (str : string) : sokolevel :=
  match (from_string_aux str.to_list) with
  | none := default sokolevel
  | some (_, ff) := default sokolevel
  | some (level, tt) := level
  end

--                              _
--    _____  ___ __   ___  _ __| |_
--   / _ \ \/ / '_ \ / _ \| '__| __|
--  |  __/>  <| |_) | (_) | |  | |_
--   \___/_/\_\ .__/ \___/|_|   \__|
--            |_|

def square_to_char : bool → bool → bool → bool → char
| tt ff ff ff := ' '
| ff ff ff ff := '#'
| tt ff tt ff := '.'
| tt tt ff ff := '$'
| tt tt tt ff := '*'
| tt ff ff tt := '@'
| tt ff tt tt := '+'
| _ _ _ _ := '?'

def to_string_aux1 (str : list char) : list (bool × bool × bool × bool) → list char
| [] := str
| ((av,box,stor,sk)::t) := (square_to_char av box stor sk)::(to_string_aux1 t)

def to_string_aux2 : list2d (bool × bool × bool × bool) → list char
| [] := []
| (h::t) := to_string_aux1 ('\n'::(to_string_aux2 t)) h

instance : has_to_string sokolevel := ⟨λ s,
  list.as_string (to_string_aux2
  (s.avail.dfzip2d (s.ini.boxes.dfzip2d (s.goal.boxes.dfzip2d (list2d.set2d true [] s.ini.storekeeper)))))
⟩

instance : has_repr sokolevel
:= ⟨λ lev, (string.append (string.append "from_string \"" (to_string lev)) "\"")⟩

meta def to_html (lev : sokolevel) : widget.html empty
  := sokowidget.build_table lev.avail lev.avail lev.ini.boxes lev.goal.boxes (bset2d.from_index lev.ini.storekeeper)

--       _                 _ _  __ _           _   _
--   ___(_)_ __ ___  _ __ | (_)/ _(_) ___ __ _| |_(_) ___  _ __
--  / __| | '_ ` _ \| '_ \| | | |_| |/ __/ _` | __| |/ _ \| '_ \
--  \__ \ | | | | | | |_) | | |  _| | (_| (_| | |_| | (_) | | | |
--  |___/_|_| |_| |_| .__/|_|_|_| |_|\___\__,_|\__|_|\___/|_| |_|
--                  |_|

meta def soko_show : tactic unit :=
do
  `(sokolevel.solvable %%lev_e) ← tactic.target,
  lev ← tactic.eval_expr sokolevel lev_e,
  tactic.trace (to_string lev)

meta def soko_simp_root (e : expr) : tactic unit :=
do
  soko ← tactic.eval_expr sokolevel e,
  tactic.trace (to_string soko),
  let p : pexpr := ``(%%e = sokolevel.mk
    %%soko.avail (sokostate.mk %%soko.ini.boxes %%soko.ini.storekeeper) (boxes_only.mk %%soko.goal.boxes)),
  eq ← tactic.to_expr p,
  name ← tactic.get_unused_name,
  H ← tactic.assert name eq,
  tactic.reflexivity,
  tactic.rewrite_target H,
  tactic.clear H

meta def soko_simp : tactic unit :=
do
  `(sokolevel.solvable %%lev) ← tactic.target,
  soko_simp_root lev

meta def soko_check_depth : expr → ℕ → tactic unit
| _ 0 := return ()
| e (n+1) := do
  `(sokolevel.move %%lev _) ← return e,
  soko_check_depth lev n

meta def soko_simp_if_deep : tactic unit :=
do
  `(sokolevel.solvable %%lev) ← tactic.target,
  tactic.try ((soko_check_depth lev 20) >> (soko_simp_root lev))

meta def solve_up : tactic unit
:= tactic.apply `(solvable.move direction.up) >> soko_simp_if_deep
meta def solve_down : tactic unit
:= tactic.apply `(solvable.move direction.down) >> soko_simp_if_deep
meta def solve_left : tactic unit
:= tactic.apply `(solvable.move direction.left) >> soko_simp_if_deep
meta def solve_right : tactic unit
:= tactic.apply `(solvable.move direction.right) >> soko_simp_if_deep
meta def solve_finish : tactic unit
:= tactic.apply `(solvable.triv) >> tactic.apply `(@of_as_true) >> tactic.triv

end sokolevel


-- Copyright (c) Microsoft Corporation. All rights reserved.
-- Licensed under the MIT license.

import ..smtexpr
import ..smtcompile
import ..bitvector
import ..verifyopt
import .spec
import .lemmas
import .irstate
import .freevar
import .equiv
import .openprog
import ..irsem_exec
import smt2.syntax
import system.io
import init.meta.tactic
import init.meta.interactive

namespace spec

open opt
open irsem

set_option pp.proofs true
lemma check_val_exec_spec_prf : check_val_exec_spec
:= begin
  unfold check_val_exec_spec,
  intros,
  cases vsrc with szsrc vsrc psrc,
  cases vtgt with sztgt vtgt ptgt,
  unfold check_val at H,
  have HSZ:decidable (sztgt = szsrc), apply_instance,
  cases HSZ,
  { rw dif_neg at H, cases H, apply neq_symm, assumption },
  {
    rw dif_pos at H,
    rw valty_rwsize_exec sztgt szsrc,
    injection H with H,
    generalize H0: (vsrc =_{irsem_exec.boolty}
      cast (check_val._match_1._proof_1
        irsem_exec szsrc sztgt (eq.symm HSZ)) vtgt) = t,
    rw H0 at H,
    unfold has_not.not at H,
    unfold has_or.or at H,
    unfold has_and.and at H,

    cases psrc,
    {
      apply val_refines.poison_intty, refl
    },
    {
      apply val_refines.concrete_intty,
      { refl },
      { cases ptgt,
        {
          cases t, cases H, cases H
        },
        refl },
      {
        cases t,
        { cases ptgt, cases H, cases H },
        {
          unfold has_eq.eq at H0,
          unfold has_comp.eq at H0,
          unfold bitvector.eq at H0,
          simp at H0, rw H0
        }
      }
    },
    any_goals { assumption }
  }
end

set_option pp.proofs true
lemma check_val_replace: ∀ vssrc vstgt (η:freevar.env),
  η⟦opt.check_val irsem_smt vssrc vstgt⟧' = opt.check_val irsem_smt (η⟦vssrc⟧) (η⟦vstgt⟧)
:= begin
  intros,
  cases vssrc, cases vstgt,
  unfold opt.check_val,
  unfold freevar.env.replace_valty,
  have HSZ: decidable (vssrc_sz = vstgt_sz), apply_instance,
  cases HSZ,
  {
    unfold check_val._match_1,
    rw dif_neg HSZ,
    rw dif_neg HSZ
  },
  {
    unfold check_val._match_1,
    rw dif_pos HSZ,
    rw dif_pos HSZ,
    unfold apply,
    unfold_coes,
    unfold id,
    rw env.replace_sb_or,
    rw env.replace_sb_not,
    rw env.replace_sb_and,
    unfold has_eq.eq, unfold has_comp.eq,
    rw env.replace_sb_eqbv,
    rw env.replace_sbv_cast, rw HSZ
  }
end

lemma check_val_some: ∀ {vssrc vstgt vesrc vetgt sres eres}
    (HEQS:equals_size vssrc vesrc = tt)
    (HEQT:equals_size vstgt vetgt = tt)
    (HS:sres = opt.check_val irsem_smt vssrc vstgt)
    (HE:eres = opt.check_val irsem_exec vesrc vetgt),
  none_or_some sres eres (λ s e, true)
:= begin
  intros,
  cases vssrc, cases vesrc, cases vstgt, cases vetgt,
  unfold equals_size at HEQS, simp at HEQS,
  unfold equals_size at HEQT, simp at HEQT,
  unfold check_val at HS,
  unfold check_val at HE,

  have HSZ': decidable (vssrc_sz = vstgt_sz), apply_instance,
  cases HSZ',
  {
    have HSZ'': vesrc_sz ≠ vetgt_sz,
    { rw ← HEQS, rw ← HEQT, assumption },
    rw dif_neg HSZ' at HS,
    rw dif_neg HSZ'' at HE,
    unfold none_or_some, left, split; assumption
  },
  {
    have HSZ'': vesrc_sz = vetgt_sz,
    { rw ← HEQS, rw ← HEQT, assumption },
    rw dif_pos HSZ' at HS,
    rw dif_pos HSZ'' at HE,
    cases sres, cases HS, injection HS,
    cases eres, cases HE, injection HE,
    unfold none_or_some, right, apply exists.intro,
    apply exists.intro, split, refl, split, refl, constructor
  }
end

lemma check_val_equiv: ∀ {vssrc vstgt vesrc vetgt sres eres}
    (HEQS:val_equiv vssrc vesrc)
    (HEQT:val_equiv vstgt vetgt)
    (HS:sres = opt.check_val irsem_smt vssrc vstgt)
    (HE:eres = opt.check_val irsem_exec vesrc vetgt),
  none_or_some sres eres (λ s e, b_equiv s e)
:= begin
  intros,
  cases vssrc, cases vstgt,
  cases vesrc, cases vetgt, -- make ivals
  have HSSZEQ := val_equiv_eqsize HEQS,
  have HTSZEQ := val_equiv_eqsize HEQT,
  unfold opt.check_val at *,

  have HSZ': decidable (vssrc_sz = vstgt_sz), apply_instance,
  cases HSZ',
  {
    have HSZ'': vesrc_sz ≠ vetgt_sz,
    {
      rw ← HSSZEQ, rw ← HTSZEQ, assumption
    },
    rw dif_neg HSZ' at HS,
    rw dif_neg HSZ'' at HE,
    rw HS, rw HE, unfold none_or_some, left, split;refl
  },
  {
    have HSZ'': vesrc_sz = vetgt_sz,
    { rw [← HSSZEQ, ← HTSZEQ], rw HSZ' },
    rw dif_pos HSZ' at HS,
    rw dif_pos HSZ'' at HE,
    cases sres, cases HS,
    cases eres, cases HE,
    injection HS with HS, injection HE with HE,
    unfold none_or_some,
    right,
    existsi sres, existsi eres,
    split, refl, split, refl,
    rw [HS, HE],

    cases HEQS,
    { -- source is poison.
      cases HEQS_a with HEQS_a,
      apply b_equiv.or1,
      { apply b_equiv.not, unfold_coes, unfold id, assumption },
      { intros H, rw HEQS_a_2 at H, cases H }
    },
    { -- source is concrete value,
      cases HEQS_a,
      cases HEQT,
      { -- target is poison
        apply b_equiv.or1, apply b_equiv.not, assumption,
        intros, apply b_equiv.and1, cases HEQT_a, assumption,
        intros, rw HEQT_a_2 at a_1, cases a_1
      },
      {
        cases HEQT_a,
        apply b_equiv.or1,
        { apply b_equiv.not, assumption },
        { intros, apply b_equiv.and1, assumption, intros,
          apply b_equiv.eq, assumption, --apply HEQT_a_1, apply bv_equiv.cast,
          apply bv_equiv_cast, assumption, rw HSZ''
        }
      }
    }
  }
end

universe u
lemma check_single_reg0_replace: ∀ psrc ptgt root ss0 (η:freevar.env),
    η⟦opt.check_single_reg0 irsem_smt psrc ptgt root ss0⟧' =
        opt.check_single_reg0 irsem_smt psrc ptgt root (η⟦ss0⟧)
:= begin
  intros,
  unfold opt.check_single_reg0,
  generalize HSS: irsem.bigstep irsem_smt ss0 psrc = oss,
  generalize HSS': irsem.bigstep irsem_smt ss0 ptgt = oss',
  rw ← bigstep_replace,
  rw ← bigstep_replace,
  rw HSS, rw HSS',
  unfold has_bind.bind,
  cases oss; cases oss';unfold option.bind,
  generalize HSV: irstate.getreg irsem_smt oss root = ovs,
  generalize HSV': irstate.getreg irsem_smt oss' root = ovs',
  rw getreg_replace HSV,
  rw getreg_replace HSV',
  cases ovs; cases ovs'; unfold option.bind,
  rw ← check_val_replace,
  generalize HCV: check_val irsem_smt ovs ovs' = ocv,
  cases ocv;unfold option.bind,
  unfold return, unfold pure,
  unfold apply,
  rw env.replace_sb_or,
  rw env.replace_sb_not,
  rw replace_getub,
  rw env.replace_sb_and,
  rw replace_getub
end

lemma check_single_reg0_equiv: ∀ psrc ptgt root ss0 se0 sres eres
    (HEQ:irstate_equiv ss0 se0)
    (HS:sres = opt.check_single_reg0 irsem_smt psrc ptgt root ss0)
    (HE:eres = opt.check_single_reg0 irsem_exec psrc ptgt root se0)
    (HEQS:irstate_equiv ss0 se0),
  none_or_some sres eres (λ s e, b_equiv s e)
:= begin
  intros,
  unfold opt.check_single_reg0 at *,
  unfold has_bind.bind at HS,
  unfold has_bind.bind at HE,
  generalize HSS: irsem.bigstep irsem_smt ss0 psrc = oss,
  generalize HSS': irsem.bigstep irsem_smt ss0 ptgt = oss',
  generalize HSE: irsem.bigstep irsem_exec se0 psrc = ose,
  generalize HSE': irsem.bigstep irsem_exec se0 ptgt = ose',
  rw [HSS, HSS'] at HS,
  rw [HSE, HSE'] at HE,
  have HPTGT_ENTANGLED: none_or_some oss' ose' (λ ss' se', irstate_equiv ss' se'),
  {
    apply bigstep_both_equiv,
    assumption, rw HSS', rw HSE'
  },
  have HPSRC_ENTANGLED: none_or_some oss ose (λ ss' se', irstate_equiv ss' se'),
  {
    apply bigstep_both_equiv,
    assumption, rw HSS, rw HSE
  },
  cases oss; cases oss';
  cases ose; cases ose',
  any_goals { unfold option.bind at HS, unfold option.bind at HE,
              rw HS, rw HE, unfold none_or_some, left, split; refl },
  any_goals {
    exfalso, apply none_or_some_false2, apply HPSRC_ENTANGLED
  },
  any_goals {
    exfalso, apply none_or_some_false1, apply HPSRC_ENTANGLED
  },
  any_goals {
    exfalso, apply none_or_some_false2, apply HPTGT_ENTANGLED
  },
  any_goals {
    exfalso, apply none_or_some_false1, apply HPTGT_ENTANGLED
  },
  rw none_or_some_apply at HPSRC_ENTANGLED,
  rw none_or_some_apply at HPTGT_ENTANGLED,
  unfold option.bind at HS,
  unfold option.bind at HE,
  generalize HSV: irstate.getreg irsem_smt oss root = ovs,
  generalize HSV': irstate.getreg irsem_smt oss' root = ovs',
  generalize HEV: irstate.getreg irsem_exec ose root = ove,
  generalize HEV': irstate.getreg irsem_exec ose' root = ove',
  have HSRCVEQ := irstate.getreg_equiv HPSRC_ENTANGLED (eq.symm HSV) (eq.symm HEV),
  have HTGTVEQ := irstate.getreg_equiv HPTGT_ENTANGLED (eq.symm HSV') (eq.symm HEV'),
  rw [HSV, HSV'] at HS,
  rw [HEV, HEV'] at HE,
  cases ovs; cases ovs'; cases ove; cases ove',
  any_goals { unfold option.bind at HS, unfold option.bind at HE,
              rw HS, rw HE, unfold none_or_some, left, split; refl },
  any_goals {
    exfalso, apply none_or_some_false2, apply HSRCVEQ
  },
  any_goals {
    exfalso, apply none_or_some_false1, apply HSRCVEQ
  },
  any_goals {
    exfalso, apply none_or_some_false2, apply HTGTVEQ
  },
  any_goals {
    exfalso, apply none_or_some_false1, apply HTGTVEQ
  },
  rw none_or_some_apply at HSRCVEQ,
  rw none_or_some_apply at HTGTVEQ,
  unfold option.bind at HS,
  unfold option.bind at HE,
  cases HSRCVEQ with HSRCSZEQ HSRCVEQ,
  cases HTGTVEQ with HTGTSZEQ HTGTVEQ,
  generalize HSCV: check_val irsem_smt ovs ovs' = scv,
  generalize HECV: check_val irsem_exec ove ove' = ecv,
  rw HSCV at HS, rw HECV at HE,
  have HCVRET := check_val_some HSRCSZEQ HTGTSZEQ (eq.symm HSCV) (eq.symm HECV),
  cases scv ; cases ecv,
  { unfold option.bind at HS, unfold option.bind at HE,
    rw HS, rw HE, unfold none_or_some, left, split; refl },
  {
    unfold option.bind at HS,
    unfold option.bind at HE,
    exfalso, apply none_or_some_false2,
    assumption
  },
  {
    unfold option.bind at HS,
    unfold option.bind at HE,
    exfalso, apply none_or_some_false1,
    assumption
  },
  unfold option.bind at HS,
  unfold option.bind at HE,
  unfold return at *, unfold pure at *,
  unfold none_or_some, right,
  apply exists.intro,
  apply exists.intro,
  split, assumption,
  split, assumption,
  generalize HUB: (~irstate.getub irsem_exec ose) = ub,
  cases ub,
  {
    apply b_equiv.or1,
    {
      rw ← HUB, apply b_equiv.not,
      apply irstate.getub_equiv,
      apply HPSRC_ENTANGLED,
      refl, refl
    },
    {
      generalize HUB': irstate.getub irsem_exec ose' = ub',
      cases ub',
      {
        intros, apply b_equiv.and1,
        rw ← HUB',
        apply irstate.getub_equiv,
        apply HPTGT_ENTANGLED, refl, rw HUB', intros Q, cases Q
      },
      {
        intros, apply b_equiv.and1,
        apply irstate.getub_equiv,
        apply HPTGT_ENTANGLED, refl, rw HUB',
        have HNOUB: has_no_ub ose = tt,
        {
          cases ose, unfold irstate.getub at HUB,
          simp at HUB, unfold has_not.not at HUB,
          cases ose_fst, cases HUB, refl
        },
        have HNOUB': has_no_ub ose' = tt,
        {
          cases ose', unfold irstate.getub at HUB',
          simp at HUB',
          cases ose'_fst, cases HUB', refl
        },
        have HV := HSRCVEQ HNOUB,
        have HV' := HTGTVEQ HNOUB',
        have HCVRET := check_val_equiv HV HV' (eq.symm HSCV) (eq.symm HECV),
        rw none_or_some_apply at HCVRET,
        intros, assumption
      }
    }
  },
  {
    apply b_equiv.or1,
    {
      rw ← HUB, apply b_equiv.not, apply irstate.getub_equiv, apply HPSRC_ENTANGLED,
      refl, refl
    },
    {
      intros Q, cases Q
    }
  }
end

theorem refines_single_reg_correct_prf: refines_single_reg_correct
:= begin
  unfold refines_single_reg_correct,
  intros,
  unfold root_refines_smt,
  intros,
  unfold root_refines_finalstate,
  intros,
  generalize HEREF: check_single_reg0 irsem_exec psrc ptgt root se0 = oeb,
  have HSREF': some (η⟦sb⟧) = check_single_reg0 irsem_smt psrc ptgt root (η⟦ss0⟧),
  {
    rw ← check_single_reg0_replace,
    rw ← HSREF
  },
  have HCHKEQ: none_or_some (some (η⟦sb⟧)) oeb (λ s e, b_equiv s e),
  {
    apply check_single_reg0_equiv,
    unfold encode at a, apply a,
    rw HSREF', rw HEREF, unfold encode at a, apply a
  },
  cases oeb with eb,
  { exfalso, apply none_or_some_false1, assumption },
  rw none_or_some_apply at HCHKEQ,
  have HCHKTT := HEQ η eb HCHKEQ,

  unfold root_refines,
  unfold check_single_reg0 at HEREF,
  rw [← a_1, ← a_2] at HEREF,
  unfold has_bind.bind at HEREF,
  unfold option.bind at HEREF,
  generalize HV: (irstate.getreg irsem_exec se root) = ov,
  generalize HV': (irstate.getreg irsem_exec se' root) = ov',
  rw [HV, HV'] at HEREF,
  cases ov with v; cases ov' with v',
  any_goals { unfold option.bind at HEREF, cases HEREF, done },
  unfold option.bind at HEREF,

  generalize HCV: check_val irsem_exec v v' = ocv,
  rw HCV at HEREF,
  cases ocv with cv,
  any_goals { unfold option.bind at HEREF, cases HEREF, done },
  unfold option.bind at HEREF,
  injection HEREF with HEREF,
  subst HCHKTT,
  split,
  {
    generalize HUBSRC : irstate.getub irsem_exec se = ubsrc,
    generalize HUBTGT : irstate.getub irsem_exec se' = ubtgt,
    rw HUBSRC at HEREF,
    rw HUBTGT at HEREF,
    cases ubsrc,
    {
      apply ub_refines.ub,
      rw HUBSRC, refl
    },
    {
      apply ub_refines.noub, rw HUBSRC, refl,
      cases ubtgt,
      { cases cv; cases HEREF },
      { rw HUBTGT, refl }
    }
  },
  {
    intros HUBTRUE,
    existsi v, existsi v',
    split, refl, split, refl,
    apply check_val_exec_spec_prf,
    rw HUBTRUE at HEREF,
    cases (irstate.getub irsem_exec se'); cases cv,
    any_goals { cases HEREF, done },
    assumption
  }
end

end spec


/-
Copyright (c) 2014 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
import Mathlib.PrePort
import Mathlib.Lean3Lib.init.logic
import Mathlib.Lean3Lib.init.control.monad
import Mathlib.Lean3Lib.init.control.alternative

universes u v u_1 u_2

namespace Mathlib

namespace option

def to_monad {m : Type → Type} [Monad m] [alternative m] {A : Type} : Option A → m A :=
  sorry

def get_or_else {α : Type u} : Option α → α → α :=
  sorry

def is_some {α : Type u} : Option α → Bool :=
  sorry

def is_none {α : Type u} : Option α → Bool :=
  sorry

def get {α : Type u} {o : Option α} : ↥(is_some o) → α :=
  sorry

def rhoare {α : Type u} : Bool → α → Option α :=
  sorry

def lhoare {α : Type u} : α → Option α → α :=
  sorry

infixr:1 "|>" => Mathlib.option.rhoare

infixr:1 "<|" => Mathlib.option.lhoare

protected def bind {α : Type u} {β : Type v} : Option α → (α → Option β) → Option β :=
  sorry

protected def map {α : Type u_1} {β : Type u_2} (f : α → β) (o : Option α) : Option β :=
  option.bind o (some ∘ f)

theorem map_id {α : Type u_1} : option.map id = id := sorry

protected instance monad : Monad Option :=
  { toApplicative :=
      { toFunctor := { map := option.map, mapConst := fun (α β : Type u_1) => option.map ∘ function.const β },
        toPure := { pure := some },
        toSeq :=
          { seq :=
              fun (α β : Type u_1) (f : Option (α → β)) (x : Option α) =>
                option.bind f fun (_x : α → β) => option.map _x x },
        toSeqLeft :=
          { seqLeft :=
              fun (α β : Type u_1) (a : Option α) (b : Option β) =>
                (fun (α β : Type u_1) (f : Option (α → β)) (x : Option α) =>
                    option.bind f fun (_x : α → β) => option.map _x x)
                  β α (option.map (function.const β) a) b },
        toSeqRight :=
          { seqRight :=
              fun (α β : Type u_1) (a : Option α) (b : Option β) =>
                (fun (α β : Type u_1) (f : Option (α → β)) (x : Option α) =>
                    option.bind f fun (_x : α → β) => option.map _x x)
                  β β (option.map (function.const α id) a) b } },
    toBind := { bind := option.bind } }

protected def orelse {α : Type u} : Option α → Option α → Option α :=
  sorry

protected instance alternative : alternative Option :=
  alternative.mk none

end option

protected instance option.inhabited (α : Type u) : Inhabited (Option α) :=
  { default := none }

protected instance option.decidable_eq {α : Type u} [d : DecidableEq α] : DecidableEq (Option α) :=
  sorry


/-
Copyright (c) 2022 Henrik Böving. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Henrik Böving
-/
import Lean.Attributes
import Lean.Environment
import Lean.Meta.Basic
import Lean.Compiler.LCNF.CompilerM

namespace Lean.Compiler.LCNF

def Phase.toNat : Phase → Nat
  | .base => 0
  | .mono => 1
  | .impure => 2

instance : LT Phase where
  lt l r := l.toNat < r.toNat

instance : LE Phase where
  le l r := l.toNat ≤ r.toNat

instance {p1 p2 : Phase} : Decidable (p1 < p2) := Nat.decLt p1.toNat p2.toNat
instance {p1 p2 : Phase} : Decidable (p1 ≤ p2) := Nat.decLe p1.toNat p2.toNat

@[simp] theorem Phase.le_refl (p : Phase) : p ≤ p := by
  cases p <;> decide

/--
A single compiler `Pass`, consisting of the actual pass function operating
on the `Decl`s as well as meta information.
-/
structure Pass where
  /--
  Which occurrence of the pass in the pipeline this is.
  Some passes, like simp, can occur multiple times in the pipeline.
  For most passes this value does not matter.
  -/
  occurrence : Nat := 0
  /--
  Which phase this `Pass` is supposed to run in
  -/
  phase : Phase
  /--
  Resulting phase.
  -/
  phaseOut : Phase := phase
  phaseInv : phaseOut ≥ phase := by simp
  /--
  The name of the `Pass`
  -/
  name : Name
  /--
  The actual pass function, operating on the `Decl`s.
  -/
  run : Array Decl → CompilerM (Array Decl)

instance : Inhabited Pass where
  default := { phase := .base, name := default, run := fun decls => return decls }

/--
Can be used to install, remove, replace etc. passes by tagging a declaration
of type `PassInstaller` with the `cpass` attribute.
-/
structure PassInstaller where
  /--
  When the installer is run this function will receive a list of all
  current `Pass`es and return a new one, this can modify the list (and
  the `Pass`es contained within) in any way it wants.
  -/
  install : Array Pass → CoreM (Array Pass)
  deriving Inhabited

/--
The `PassManager` used to store all `Pass`es that will be run within
pipeline.
-/
structure PassManager where
  passes : Array Pass
  deriving Inhabited

instance : ToString Phase where
  toString
    | .base => "base"
    | .mono => "mono"
    | .impure => "impure"

namespace Pass

def mkPerDeclaration (name : Name) (run : Decl → CompilerM Decl) (phase : Phase) (occurrence : Nat := 0) : Pass where
  occurrence := occurrence
  phase := phase
  name := name
  run := fun xs => xs.mapM run

end Pass

namespace PassManager

def validate (manager : PassManager) : CoreM Unit := do
  let mut current := .base
  for pass in manager.passes do
    if ¬(current ≤ pass.phase) then
      throwError s!"{pass.name} has phase {pass.phase} but should at least have {current}"
    current := pass.phase

def findHighestOccurrence (targetName : Name) (passes : Array Pass) : CoreM Nat := do
  let mut highest := none
  for pass in passes do
      if pass.name == targetName then
        highest := some pass.occurrence
  let some val := highest | throwError s!"Could not find any occurrence of {targetName}"
  return val

end PassManager

namespace PassInstaller

def installAtEnd (p : Pass) : PassInstaller where
  install passes := return passes.push p

def append (passesNew : Array Pass) : PassInstaller where
  install passes := return passes ++ passesNew

def withEachOccurrence (targetName : Name) (f : Nat → PassInstaller) : PassInstaller where
  install passes := do
    let highestOccurrence ← PassManager.findHighestOccurrence targetName passes
    let mut passes := passes
    for occurrence in [0:highestOccurrence+1] do
      passes ← f occurrence |>.install passes
    return passes

def installAfter (targetName : Name) (p : Pass → Pass) (occurrence : Nat := 0) : PassInstaller where
  install passes :=
    if let some idx := passes.findIdx? (fun p => p.name == targetName && p.occurrence == occurrence) then
      let passUnderTest := passes[idx]!
      return passes.insertAt! (idx + 1) (p passUnderTest)
    else
      throwError s!"Tried to insert pass after {targetName}, occurrence {occurrence} but {targetName} is not in the pass list"

def installAfterEach (targetName : Name) (p : Pass → Pass) : PassInstaller :=
    withEachOccurrence targetName (installAfter targetName p ·)

def installBefore (targetName : Name) (p : Pass → Pass) (occurrence : Nat := 0): PassInstaller where
  install passes :=
    if let some idx := passes.findIdx? (fun p => p.name == targetName && p.occurrence == occurrence) then
      let passUnderTest := passes[idx]!
      return passes.insertAt! idx (p passUnderTest)
    else
      throwError s!"Tried to insert pass after {targetName}, occurrence {occurrence} but {targetName} is not in the pass list"

def installBeforeEachOccurrence (targetName : Name) (p : Pass → Pass) : PassInstaller :=
    withEachOccurrence targetName (installBefore targetName p ·)

def replacePass (targetName : Name) (p : Pass → Pass) (occurrence : Nat := 0) : PassInstaller where
  install passes := do
    let some idx := passes.findIdx? (fun p => p.name == targetName && p.occurrence == occurrence) | throwError s!"Tried to replace {targetName}, occurrence {occurrence} but {targetName} is not in the pass list"
    let target := passes[idx]!
    let replacement := p target
    return passes.set! idx replacement

def replaceEachOccurrence (targetName : Name) (p : Pass → Pass) : PassInstaller :=
    withEachOccurrence targetName (replacePass targetName p ·)

def run (manager : PassManager) (installer : PassInstaller) : CoreM PassManager := do
  return { manager with passes := (← installer.install manager.passes) }

private unsafe def getPassInstallerUnsafe (declName : Name) : CoreM PassInstaller := do
  ofExcept <| (← getEnv).evalConstCheck PassInstaller (← getOptions) ``PassInstaller declName

@[implemented_by getPassInstallerUnsafe]
private opaque getPassInstaller (declName : Name) : CoreM PassInstaller

def runFromDecl (manager : PassManager) (declName : Name) : CoreM PassManager := do
  let installer ← getPassInstaller declName
  let newState ← installer.run manager
  newState.validate
  return newState

end PassInstaller

end Lean.Compiler.LCNF


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

! This file was ported from Lean 3 source module init.meta.smt.congruence_closure
! leanprover-community/mathlib commit 9eae65f7144bcc692858b9dadf2e48181f4270b9
! Please do not edit these lines, except to modify the commit id
! if you have ported upstream changes.
-/
prelude
import Leanbin.Init.Meta.InteractiveBase
import Leanbin.Init.Meta.Tactic
import Leanbin.Init.Meta.SetGetOptionTactics

structure CcConfig where
  -- If tt, congruence closure will treat implicit instance arguments as constants.
  ignoreInstances : Bool := true
  -- If tt, congruence closure modulo AC.
  ac : Bool := true
  /- If ho_fns is (some fns), then full (and more expensive) support for higher-order functions is
     *only* considered for the functions in fns and local functions. The performance overhead is described in the paper
     "Congruence Closure in Intensional Type Theory". If ho_fns is none, then full support is provided
     for *all* constants. -/
  hoFns : Option (List Name) := none
  -- If true, then use excluded middle
  em : Bool := true
#align cc_config CcConfig

/-- Congruence closure state.
This may be considered to be a set of expressions and an equivalence class over this set.
The equivalence class is generated by the equational rules that are added to the cc_state and congruence,
that is, if `a = b` then `f(a) = f(b)` and so on.
 -/
unsafe axiom cc_state : Type
#align cc_state cc_state

unsafe axiom cc_state.mk_core : CcConfig → cc_state
#align cc_state.mk_core cc_state.mk_core

/-- Create a congruence closure state object using the hypotheses in the current goal. -/
unsafe axiom cc_state.mk_using_hs_core : CcConfig → tactic cc_state
#align cc_state.mk_using_hs_core cc_state.mk_using_hs_core

/-- Get the next element in the equivalence class.
Note that if the given expr e is not in the graph then it will just return e. -/
unsafe axiom cc_state.next : cc_state → expr → expr
#align cc_state.next cc_state.next

/-- Returns the root expression for each equivalence class in the graph.
If the bool argument is set to true then it only returns roots of non-singleton classes. -/
unsafe axiom cc_state.roots_core : cc_state → Bool → List expr
#align cc_state.roots_core cc_state.roots_core

/-- Get the root representative of the given expression. -/
unsafe axiom cc_state.root : cc_state → expr → expr
#align cc_state.root cc_state.root

/--
"Modification Time". The field m_mt is used to implement the mod-time optimization introduce by the Simplify theorem prover.
The basic idea is to introduce a counter gmt that records the number of heuristic instantiation that have
occurred in the current branch. It is incremented after each round of heuristic instantiation.
The field m_mt records the last time any proper descendant of of thie entry was involved in a merge. -/
unsafe axiom cc_state.mt : cc_state → expr → Nat
#align cc_state.mt cc_state.mt

/-- "Global Modification Time". gmt is a number stored on the cc_state,
it is compared with the modification time of a cc_entry in e-matching. See `cc_state.mt`. -/
unsafe axiom cc_state.gmt : cc_state → Nat
#align cc_state.gmt cc_state.gmt

/-- Increment the Global Modification time. -/
unsafe axiom cc_state.inc_gmt : cc_state → cc_state
#align cc_state.inc_gmt cc_state.inc_gmt

/-- Check if `e` is the root of the congruence class. -/
unsafe axiom cc_state.is_cg_root : cc_state → expr → Bool
#align cc_state.is_cg_root cc_state.is_cg_root

/-- Pretty print the entry associated with the given expression. -/
unsafe axiom cc_state.pp_eqc : cc_state → expr → tactic format
#align cc_state.pp_eqc cc_state.pp_eqc

/-- Pretty print the entire cc graph.
If the bool argument is set to true then singleton equivalence classes will be omitted. -/
unsafe axiom cc_state.pp_core : cc_state → Bool → tactic format
#align cc_state.pp_core cc_state.pp_core

/-- Add the given expression to the graph. -/
unsafe axiom cc_state.internalize : cc_state → expr → tactic cc_state
#align cc_state.internalize cc_state.internalize

/-- Add the given proof term as a new rule.
The proof term p must be an `eq _ _`, `heq _ _`, `iff _ _`, or a negation of these. -/
unsafe axiom cc_state.add : cc_state → expr → tactic cc_state
#align cc_state.add cc_state.add

/-- Check whether two expressions are in the same equivalence class. -/
unsafe axiom cc_state.is_eqv : cc_state → expr → expr → tactic Bool
#align cc_state.is_eqv cc_state.is_eqv

/-- Check whether two expressions are not in the same equivalence class. -/
unsafe axiom cc_state.is_not_eqv : cc_state → expr → expr → tactic Bool
#align cc_state.is_not_eqv cc_state.is_not_eqv

/-- Returns a proof term that the given terms are equivalent in the given cc_state-/
unsafe axiom cc_state.eqv_proof : cc_state → expr → expr → tactic expr
#align cc_state.eqv_proof cc_state.eqv_proof

/--
Returns true if the cc_state is inconsistent. For example if it had both `a = b` and `a ≠ b` in it.-/
unsafe axiom cc_state.inconsistent : cc_state → Bool
#align cc_state.inconsistent cc_state.inconsistent

/-- `proof_for cc e` constructs a proof for e if it is equivalent to true in cc_state -/
unsafe axiom cc_state.proof_for : cc_state → expr → tactic expr
#align cc_state.proof_for cc_state.proof_for

/-- `refutation_for cc e` constructs a proof for `not e` if it is equivalent to false in cc_state -/
unsafe axiom cc_state.refutation_for : cc_state → expr → tactic expr
#align cc_state.refutation_for cc_state.refutation_for

/-- If the given state is inconsistent, return a proof for false. Otherwise fail. -/
unsafe axiom cc_state.proof_for_false : cc_state → tactic expr
#align cc_state.proof_for_false cc_state.proof_for_false

namespace CcState

unsafe def mk : cc_state :=
  cc_state.mk_core { }
#align cc_state.mk cc_state.mk

unsafe def mk_using_hs : tactic cc_state :=
  cc_state.mk_using_hs_core { }
#align cc_state.mk_using_hs cc_state.mk_using_hs

unsafe def roots (s : cc_state) : List expr :=
  cc_state.roots_core s true
#align cc_state.roots cc_state.roots

unsafe instance : has_to_tactic_format cc_state :=
  ⟨fun s => cc_state.pp_core s true⟩

unsafe def eqc_of_core (s : cc_state) : expr → expr → List expr → List expr
  | e, f, r =>
    let n := s.next e
    if n = f then e :: r else eqc_of_core n f (e :: r)
#align cc_state.eqc_of_core cc_state.eqc_of_core

unsafe def eqc_of (s : cc_state) (e : expr) : List expr :=
  s.eqc_of_core e e []
#align cc_state.eqc_of cc_state.eqc_of

unsafe def in_singlenton_eqc (s : cc_state) (e : expr) : Bool :=
  s.next e = e
#align cc_state.in_singlenton_eqc cc_state.in_singlenton_eqc

unsafe def eqc_size (s : cc_state) (e : expr) : Nat :=
  (s.eqc_of e).length
#align cc_state.eqc_size cc_state.eqc_size

unsafe def fold_eqc_core {α} (s : cc_state) (f : α → expr → α) (first : expr) : expr → α → α
  | c, a =>
    let new_a := f a c
    let next := s.next c
    if next == first then new_a else fold_eqc_core next new_a
#align cc_state.fold_eqc_core cc_state.fold_eqc_core

unsafe def fold_eqc {α} (s : cc_state) (e : expr) (a : α) (f : α → expr → α) : α :=
  fold_eqc_core s f e e a
#align cc_state.fold_eqc cc_state.fold_eqc

unsafe def mfold_eqc {α} {m : Type → Type} [Monad m] (s : cc_state) (e : expr) (a : α)
    (f : α → expr → m α) : m α :=
  fold_eqc s e (return a) fun act e => do
    let a ← act
    f a e
#align cc_state.mfold_eqc cc_state.mfold_eqc

end CcState

open Tactic

unsafe def tactic.cc_core (cfg : CcConfig) : tactic Unit := do
  intros
  let s ← cc_state.mk_using_hs_core cfg
  let t ← target
  let s ← s.internalize t
  if s then do
      let pr ← s
      mk_app `false.elim [t, pr] >>= exact
    else do
      let tr ← return <| expr.const `true []
      let b ← s t tr
      if b then do
          let pr ← s t tr
          mk_app `of_eq_true [pr] >>= exact
        else do
          let dbg ← get_bool_option `trace.cc.failure ff
          if dbg then do
              let ccf ← pp s
              fail
                  f! "cc tactic failed, equivalence classes:
                    {ccf}"
            else do
              fail "cc tactic failed"
#align tactic.cc_core tactic.cc_core

unsafe def tactic.cc : tactic Unit :=
  tactic.cc_core { }
#align tactic.cc tactic.cc

unsafe def tactic.cc_dbg_core (cfg : CcConfig) : tactic Unit :=
  save_options <| set_bool_option `trace.cc.failure true >> tactic.cc_core cfg
#align tactic.cc_dbg_core tactic.cc_dbg_core

unsafe def tactic.cc_dbg : tactic Unit :=
  tactic.cc_dbg_core { }
#align tactic.cc_dbg tactic.cc_dbg

unsafe def tactic.ac_refl : tactic Unit := do
  let (lhs, rhs) ← target >>= match_eq
  let s ← return <| cc_state.mk
  let s ← s.internalize lhs
  let s ← s.internalize rhs
  let b ← s.is_eqv lhs rhs
  if b then do
      s lhs rhs >>= exact
    else do
      fail "ac_refl failed"
#align tactic.ac_refl tactic.ac_refl