(** The "symbolic" AST is the AST directly generated by the symbolic execution.
    It is very rough and meant to be extremely straightforward to build during
    the symbolic execution: we later apply transformations to generate the pure
    AST that we export. *)

open Types
open Expressions
open Values
open LlbcAst
module SymbolicExprId = Contexts.SymbolicExprId

type symbolic_expr_id = Contexts.symbolic_expr_id [@@deriving show, eq, ord]

(** "Meta"-place: a place stored as span-data.

    Whenever we need to introduce new symbolic variables, for instance because
    of symbolic expansions, we try to store a "place", which gives information
    about the origin of the values (this place information comes from assignment
    information, etc.). We later use this place information to generate
    meaningful name, to prettify the generated code. *)
type mplace =
  | PlaceLocal of real_var_binder
      (** It is important that we store the binder, and not just the variable
          id, because the most important information in a place is the name of
          the variable! *)
  | PlaceGlobal of global_decl_ref
  | PlaceProjection of mplace * projection_elem
[@@deriving show]

type call_id =
  | Fun of fn_ptr_kind * FunCallId.id
      (** A "regular" function (i.e., a function which is not a primitive
          operation) *)
  | Unop of unop
  | Binop of binop
[@@deriving show, ord]

type call = {
  call_id : call_id;
  span : Meta.span;
  ctx : (Contexts.eval_ctx[@opaque]);
      (** The context upon calling the function (after the operands have been
          evaluated). We need it to compute the translated values for shared
          borrows (we need to perform lookups). *)
  sg : fun_sig option;
      (** The non-instantiated function signature, if this is not a unop/binop.

          This is useful to retrieve the names of the inputs, to generate pretty
          names in the translation. *)
  inst_sg : inst_fun_sig option;
      (** The instantiated function signature, if this is not a unop/binop *)
  abstractions : AbsId.id list;
      (** The region abstractions introduced upon calling the function *)
  generics : generic_args;
  args : tvalue list;
  args_places : mplace option list;  (** Meta information *)
  dest : symbolic_value;
  dest_place : mplace option;  (** Meta information *)
}
[@@deriving show]

(** Meta information for expressions, not necessary for synthesis but useful to
    guide it to generate a pretty output. *)
type emeta =
  | Assignment of Contexts.eval_ctx * mplace * tvalue * mplace option
      (** We generated an assignment (destination, assigned value, src) *)
  | Snapshot of meta_id * Contexts.eval_ctx
      (** Remember an environment snapshot - this is useful to check where the
          symbolic values are, to compute proper names for instance.

          We use the meta id for debugging purposes: it allows us to use the
          option `-mark-ids` to stop the execution when this meta-information is
          introduced (this is useful, because it allows us to figure out where
          for instance the context introduced in the meta-information comes
          from). *)
[@@deriving show]

type variant_id = VariantId.id [@@deriving show]
type global_decl_id = GlobalDeclId.id [@@deriving show]
type 'a symbolic_value_id_map = 'a SymbolicValueId.Map.t [@@deriving show]
type 'a abs_id_map = 'a AbsId.Map.t [@@deriving show]
type 'a region_group_id_map = 'a RegionGroupId.Map.t [@@deriving show]

(** Ancestor for {!expr} iter visitor.

    We could make it inherit the visitor for {!Contexts.eval_ctx}, but in all
    the uses of this visitor we don't need to explore {!Contexts.eval_ctx}, so
    we make it inherit the abstraction visitors instead. *)
class ['self] iter_expr_base =
  object (self : 'self)
    inherit [_] iter_abs
    method visit_eval_ctx : 'env -> Contexts.eval_ctx -> unit = fun _ _ -> ()
    method visit_call : 'env -> call -> unit = fun _ _ -> ()
    method visit_mplace : 'env -> mplace -> unit = fun _ _ -> ()
    method visit_emeta : 'env -> emeta -> unit = fun _ _ -> ()
    method visit_abs_level : 'env -> abs_level -> unit = fun _ _ -> ()

    method visit_region_group_id_map :
        'a. ('env -> 'a -> unit) -> 'env -> 'a region_group_id_map -> unit =
      fun f env m ->
        RegionGroupId.Map.iter
          (fun id x ->
            self#visit_region_group_id env id;
            f env x)
          m

    method visit_symbolic_value_id_map :
        'a. ('env -> 'a -> unit) -> 'env -> 'a symbolic_value_id_map -> unit =
      fun f env m ->
        SymbolicValueId.Map.iter
          (fun id x ->
            self#visit_symbolic_value_id env id;
            f env x)
          m

    method visit_abs_id_map :
        'a. ('env -> 'a -> unit) -> 'env -> 'a abs_id_map -> unit =
      fun f env m ->
        AbsId.Map.iter
          (fun id x ->
            self#visit_abs_id env id;
            f env x)
          m

    method visit_symbolic_value_id_set : 'env -> symbolic_value_id_set -> unit =
      fun env s -> SymbolicValueId.Set.iter (self#visit_symbolic_value_id env) s

    method visit_symbolic_expr_id : 'env -> symbolic_expr_id -> unit =
      fun _ _ -> ()

    method visit_symbolic_expansion : 'env -> symbolic_expansion -> unit =
      fun _ _ -> ()
  end

(** **Rem.:** here, {!expression} is not at all equivalent to the expressions
    used in LLBC or in lambda-calculus: they are simply a first step towards
    lambda-calculus expressions. *)
type expr =
  | Return of (Contexts.eval_ctx[@opaque]) * tvalue option
      (** There are two cases:
          - the AST is for a forward function: the typed value should contain
            the value which was in the return variable
          - the AST is for a backward function: the typed value should be [None]

          The context is the evaluation context upon reaching the return, We
          need it to translate shared borrows to pure values (we need to be able
          to look up the shared values in the context). *)
  | Panic
  | FunCall of call * expr
  | EndAbs of (Contexts.eval_ctx[@opaque]) * abs * abs_level * expr
      (** The context is the evaluation context upon ending the abstraction,
          just after we removed the abstraction from the context.

          The integer is the level of the sub-abstraction that we ended. *)
  | EvalGlobal of global_decl_id * generic_args * symbolic_value * expr
      (** Evaluate a global to a fresh symbolic value *)
  | Assertion of (Contexts.eval_ctx[@opaque]) * bool * tvalue * expr
      (** An assertion.

          The context is the evaluation context from after evaluating the
          asserted value. It has the same purpose as for the {!Return} case.

          The boolean is the expected value (the [tvalue] should evaluate to
          this boolean for the assertion to succeed). *)
  | Expansion of mplace option * symbolic_value * expansion
      (** Expansion of a symbolic value.

          The place is "span": it gives the path to the symbolic value (if
          available) which got expanded (this path is available when the
          symbolic expansion comes from a path evaluation, during an assignment
          for instance). We use it to compute meaningful names for the variables
          we introduce, to prettify the generated code. *)
  | IntroSymbolic of
      (Contexts.eval_ctx[@opaque])
      * mplace option
      * symbolic_value
      * value_aggregate
      * expr
      (** We introduce a new symbolic value, equal to some other value.

          This is used for instance when reorganizing the environment to compute
          fixed points: we duplicate some shared symbolic values to destructure
          the shared values, in order to make the environment a bit more general
          (while losing precision of course). We also use it to introduce
          symbolic values when evaluating constant generics, or trait constants.

          The context is the evaluation context from before introducing the new
          value. It has the same purpose as for the {!Return} case. *)
  | SubstituteAbsIds of symbolic_expr_id * abs_id abs_id_map * expr
      (** We sometimes need to substitute abstraction ids to refresh them (in
          particular when doing joins), which can be a problem especially as
          some abstraction expressions refer to the abstractions through their
          ids. In order to make the translation work, we need to save those
          substitutions.

          TODO: we actually don't do this anymore because it led to issues.
          Remove? *)
  | ForwardEnd of
      ((Contexts.eval_ctx[@opaque]) * tvalue) option
      * (Contexts.eval_ctx[@opaque])
      * expr
      * expr region_group_id_map
      (** We use this delimiter to indicate at which point we switch to the
          generation of code specific to the backward function(s).

          The fields are:
          - optional: the evaluation context **after we evaluated the return
            value** with the value consumed by the return variable
          - the evaluation context at the moment we introduce the [ForwardEnd].
            We use it to translate the input values (see the comments for the
            {!Return} variant). TODO: it seems this is not necessary anymore.
          - the end of the translation for the forward function
          - a map from region group ids to expressions that give the end of the
            translation for the backward functions

          This case also handles the case where we (re-)enter a loop (once we
          enter a loop in symbolic mode, we don't get out: the loop is
          responsible for the end of the function).

          TODO: because we store the returned value, the Return case may not be
          useful anymore? *)
  | LoopContinue of
      (Contexts.eval_ctx[@opaque]) * loop_id * tvalue list * abs list
  | LoopBreak of (Contexts.eval_ctx[@opaque]) * loop_id * tvalue list * abs list
  | Loop of loop  (** Loop: call to a loop *)
  | Let of let_expr  (** A let binding. See the comments for [let_expr]. *)
  | Join of (Contexts.eval_ctx[@opaque]) * tvalue list * abs list
      (** The output of an expression which is bound (this is the result of
          matching the context at the end of a branch with the joined context).

          Ex.:
          {[
            if b {
              x = 0;
            }
            else {
              x = 1;
            }; // We need to join contexts here
            x + y;
          ]}

          The symbolic AST resulting from symbolically executing the code looks
          like this:
          {[
            Let {
              bound = (
                ExpansionBool // introduced by the if then else
                   (Join (..., [0], []), // then branch
                    Join (..., [1], [])  // else branch
                    ));
              next_expr = ...; // result of translating [x + y]
              ... }
          ]} *)
  | Meta of (emeta[@opaque]) * expr  (** Meta information *)
  | Error of Meta.span option * string
  | TargetDispatch of symbolic_value list * (string * Types.fun_decl_ref) list
      (** Multi-target dispatch: the function body dispatches to one of several
          per-target implementations.

          The first field is the list of input symbolic values (the function
          arguments to forward to each target).

          Each element of the second list is [(target_name, fun_ref)] where
          [fun_ref] identifies the per-target function to call. *)

and loop = {
  ctx : (Contexts.eval_ctx[@opaque]);
      (** The evaluation context just before the loop *)
  loop_id : loop_id;
  input_svalues : symbolic_value list;
      (** The input symbolic values, properly ordered *)
  input_abs : abs list;
      (** The input abstractions, properly ordered. Note that those are the
          abstractions from the *fixed-point*, they are not actually the
          abstractions received as input from the loop (see [input_abs_to_abs]
          instead). We store abstractions (rather than ids) because we need the
          abstractions themselves to compute the type of the continuations
          (i.e., the type of the loop) and doing this avoids a lookup. *)
  input_value_to_value : tvalue symbolic_value_id_map;
  input_abs_to_abs : abs abs_id_map;
  break_svalues : symbolic_value list;
      (** The symbolic values introduced in the break environment (those are
          output by the loop) *)
  break_abs : abs list;
      (** The abstractions introduced in the break environment (those are output
          by the loop) *)
  loop_expr : expr;  (** The symbolically executed loop body *)
  next_expr : expr;  (** The expression for *after* the loop call *)
  span : Meta.span;  (** Information about the origin of the loop body *)
}

(** A let-binding.

    We use this when joining the contexts after a branching expression (e.g., an
    [if then else]): the resulting pure code binds the "output" of the
    [if then else]. *)
and let_expr = {
  bound_expr : expr;  (** The bound expression *)
  out_svalues : symbolic_value list;
      (** The symbolic values introduced in the joined environment (those are
          output by the bound expression) *)
  out_abs : abs list;
      (** The abstractions introduced in the joined environment (those are
          output by the bound expression) *)
  next_expr : expr;  (** The expression for *after* the let *)
  span : Meta.span;  (** Information about the origin of the expression *)
}

and expansion =
  | ExpandNoBranch of symbolic_expansion * expr
      (** A symbolic expansion which doesn't generate a branching. Includes:
          - concrete expansion
          - borrow expansion *Doesn't* include:
          - expansion of ADTs with one variant *)
  | ExpandAdt of (variant_id option * symbolic_value list * expr) list
      (** ADT expansion *)
  | ExpandBool of expr * expr
      (** A boolean expansion (i.e, an [if ... then ... else ...]) *)
  | ExpandInt of integer_type * (scalar_value * expr) list * expr
      (** An integer expansion (i.e, a switch over an integer). The last
          expression is for the "otherwise" branch. *)

(* Remark: this type doesn't have to be mutually recursive with the other
   types, but it makes it easy to generate the visitors *)
and value_aggregate =
  | VaSingleValue of tvalue  (** Regular case *)
  | VaArray of tvalue list
      (** This is used when introducing array aggregates *)
  | VaCgValue of const_generic_var_id
      (** This is used when evaluating a const generic value: in the
          interpreter, we introduce a fresh symbolic value. *)
  | VaTraitConstValue of trait_ref * assoc_const_id
      (** A trait constant value *)
  | VaDiscriminant of symbolic_value  (** A discriminant read *)
  | VaDynTrait of tvalue * trait_ref
      (** A dynamic trait. This gets inserted when we convert a box of an
          element of a known type to a box of an element of type [dyn] through
          an unsized cast: we need the trait reference to perform the cast. *)
  | VaFnDef of fn_ptr  (** Function pointer of a top-level definition *)
[@@deriving
  show,
  visitors
    {
      name = "iter_expr";
      variety = "iter";
      ancestors = [ "iter_expr_base" ];
      nude = true (* Don't inherit {!VisitorsRuntime.iter} *);
      concrete = true;
      polymorphic = false;
    }]
