val sort:permute
let sort : permute = my_sortWith #int (my_compare_of_bool (<))
val sort:permute
let sort:permute =
my_sortWith #int (my_compare_of_bool ( < ))
Steel.Effect.Common.fsti.checked
Steel.Memory.fsti.checked prims.fst.checked FStar.Tactics.V2.fst.checked FStar.Tactics.CanonCommMonoidSimple.Equiv.fst.checked FStar.String.fsti.checked FStar.Squash.fsti.checked FStar.Set.fsti.checked FStar.Reflection.V2.Derived.Lemmas.fst.checked FStar.Pervasives.Native.fst.checked FStar.Pervasives.fsti.checked FStar.List.Tot.Base.fst.checked FStar.List.Tot.fst.checked FStar.Ghost.fsti.checked FStar.FunctionalExtensionality.fsti.checked FStar.Classical.fsti.checked FStar.Algebra.CommMonoid.Equiv.fst.checked
Steel.Effect.Common.fsti
Steel.Effect.Common.my_sortWith Prims.int Steel.Effect.Common.my_compare_of_bool Prims.op_LessThan
(*
   Copyright 2020 Microsoft Research

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*)

module Steel.Effect.Common
open Steel.Memory
module Mem = Steel.Memory
module FExt = FStar.FunctionalExtensionality
open FStar.Ghost

/// This module provides various predicates and functions which are common to the
/// different Steel effects.
/// It also contains the tactic responsible for frame inference through a variant of AC-unification

#set-options "--ide_id_info_off"

(* Normalization helpers *)

irreducible let framing_implicit : unit = ()
irreducible let __steel_reduce__ : unit = ()
/// An internal attribute for finer-grained normalization in framing equalities
irreducible let __inner_steel_reduce__ : unit = ()
irreducible let __reduce__ : unit = ()
irreducible let smt_fallback : unit = ()
irreducible let ite_attr : unit = ()

// Needed to avoid some logical vs prop issues during unification with no subtyping
[@@__steel_reduce__]
unfold
let true_p : prop = True

module T = FStar.Tactics.V2

let join_preserves_interp (hp:slprop) (m0:hmem hp) (m1:mem{disjoint m0 m1})
: Lemma
  (interp hp (join m0 m1))
  [SMTPat (interp hp (join m0 m1))]
= let open Steel.Memory in
  intro_emp m1;
  intro_star hp emp m0 m1;
  affine_star hp emp (join m0 m1)

(* Definition of a selector for a given slprop *)

/// A selector of type `a` for a separation logic predicate hp is a function
/// from a memory where the predicate hp holds, which returns a value of type `a`.
/// The effect GTot indicates that selectors are ghost functions, used for specification
/// and proof purposes only
let selector' (a:Type0) (hp:slprop) = hmem hp -> GTot a

/// Self-framing property for selectors
let sel_depends_only_on (#a:Type) (#hp:slprop) (sel:selector' a hp) =
  forall (m0:hmem hp) (m1:mem{disjoint m0 m1}).
    (interp_depends_only_on hp; (
    sel m0 == sel (join m0 m1)))

/// Additional property that selectors must satisfy, related to internals of
/// the Steel memory model encoding
let sel_depends_only_on_core (#a:Type) (#hp:slprop) (sel:selector' a hp) =
  forall (m0:hmem hp). sel m0 == sel (core_mem m0)

/// Full definition of a selector, as a function which satisfies the two predicates above
let selector (a:Type) (hp:slprop) : Type =
  sel:selector' a hp{sel_depends_only_on sel /\ sel_depends_only_on_core sel}

/// The basis of our selector framework: Separation logic assertions enhanced with selectors
/// Note that selectors are "optional", it is always possible to use a non-informative selector,
/// such as fun _ -> () and to rely on the standard separation logic reasoning
[@@ erasable]
noeq
type vprop' =
  { hp: slprop u#1;
    t:Type0;
    sel: selector t hp}

(* Lifting the star operator to an inductive type makes normalization
   and implementing some later functions easier *)
[@@__steel_reduce__; erasable]
noeq
type vprop =
  | VUnit : vprop' -> vprop
  | VStar: vprop -> vprop -> vprop

(* A generic lift from slprop to vprop with a non-informative selector *)

[@@ __steel_reduce__]
let to_vprop' (p:slprop) = {hp = p; t = unit; sel = fun _ -> ()}

[@@ __steel_reduce__]
unfold
let to_vprop (p:slprop) = VUnit (to_vprop' p)

/// Normalization steps for norm below.
/// All functions marked as `unfold`, or with the `__steel_reduce__` attribute will be reduced,
/// as well as some functions internal to the selector framework
unfold
let normal_steps =
   [delta_attr [`%__steel_reduce__; `%__inner_steel_reduce__];
    delta_only [`%Mkvprop'?.t; `%Mkvprop'?.hp; `%Mkvprop'?.sel;
      `%FStar.Algebra.CommMonoid.Equiv.__proj__CM__item__mult;
      `%FStar.Algebra.CommMonoid.Equiv.__proj__CM__item__unit];
    delta_qualifier ["unfold"];
    iota;zeta;primops; simplify]

/// The core normalization primitive used to simplify Verification Conditions before encoding
/// them to an SMT solver.
unfold
let normal (#a:Type) (x:a) = norm normal_steps x

/// An abbreviation for the VStar constructor, allowing to use it with infix notation
[@@ __steel_reduce__; __reduce__]
let star = VStar

/// Extracting the underlying separation logic assertion from a vprop
[@@ __steel_reduce__]
let rec hp_of (p:vprop) = match p with
  | VUnit p -> p.hp
  | VStar p1 p2 -> hp_of p1 `Mem.star` hp_of p2

/// Extracting the selector type from a vprop
[@@ __steel_reduce__]
let rec t_of (p:vprop) = match p with
  | VUnit p -> p.t
  | VStar p1 p2 -> t_of p1 * t_of p2

/// Extracting the selector from a vprop
[@@ __steel_reduce__]
let rec sel_of (p:vprop) : GTot (selector (t_of p) (hp_of p)) = match p with
  | VUnit p -> fun h -> p.sel h
  | VStar p1 p2 ->
     let sel1 = sel_of p1 in
     let sel2 = sel_of p2 in
     fun h -> (sel1 h, sel2 h)

/// Type abbreviations for separation logic pre- and postconditions of the Steel effects
type pre_t = vprop
type post_t (a:Type) = a -> vprop

/// An annotation to indicate which separation logic predicates correspond to monadic computations
/// These computations are handled in a specific manner in the framing tactic; they correspond to places where
/// the context shrinks from all local variables in the computation to variables available at the toplevel
let return_pre (p:vprop) : vprop = p

noextract
let hmem (p:vprop) = hmem (hp_of p)

/// Abstract predicate for vprop implication. Currently implemented as an implication on the underlying slprop
val can_be_split (p q:pre_t) : Type0

/// Exposing the implementation of `can_be_split` when needed for proof purposes
val reveal_can_be_split (_:unit) : Lemma
  (forall p q. can_be_split p q == Mem.slimp (hp_of p) (hp_of q))

/// A targeted version of the above
val can_be_split_interp (r r':vprop) (h:hmem r)
  : Lemma (requires can_be_split r r')
          (ensures interp (hp_of r') h)


/// A dependent version of can_be_split, to be applied to dependent postconditions
let can_be_split_forall (#a:Type) (p q:post_t a) = forall x. can_be_split (p x) (q x)

/// A version of can_be_split which is indexed by a proposition, which can be used for equalities abduction
let can_be_split_dep (p:prop) (t1 t2:pre_t) = p ==> can_be_split t1 t2

/// A dependent version of the above predicate
let can_be_split_forall_dep (#a:Type) (p:a -> prop) (t1 t2:post_t a) =
  forall (x:a). p x ==> can_be_split (t1 x) (t2 x)

(* Some lemmas about the can_be_split* predicates,
  to be used as rewriting rules for the abstract predicates *)

val can_be_split_trans (p q r:vprop)
: Lemma
  (requires p `can_be_split` q /\ q `can_be_split` r)
  (ensures p `can_be_split` r)

let can_be_split_trans_rev (p q r:vprop)
: Lemma
  (requires q `can_be_split` r /\ p `can_be_split` q)
  (ensures p `can_be_split` r)
= can_be_split_trans p q r

val can_be_split_star_l (p q:vprop)
: Lemma
  (ensures (p `star` q) `can_be_split` p)
  [SMTPat ((p `star` q) `can_be_split` p)]

val can_be_split_star_r (p q:vprop)
: Lemma
  (ensures (p `star` q) `can_be_split` q)
  [SMTPat ((p `star` q) `can_be_split` q)]

val can_be_split_refl (p:vprop)
: Lemma (p `can_be_split` p)
[SMTPat (p `can_be_split` p)]


val can_be_split_congr_l
  (p q r: vprop)
: Lemma
  (requires (p `can_be_split` q))
  (ensures ((p `star` r) `can_be_split` (q `star` r)))

val can_be_split_congr_r
  (p q r: vprop)
: Lemma
  (requires (p `can_be_split` q))
  (ensures ((r `star` p) `can_be_split` (r `star` q)))

let prop_and (p1 p2: prop) : Tot prop = p1 /\ p2

let can_be_split_forall_dep_trans_rev
  (#a: Type)
  (cond1 cond2: a -> prop)
  (p q r: post_t a)
: Lemma
  (requires (can_be_split_forall_dep cond2 q r /\ can_be_split_forall_dep cond1 p q))
  (ensures (can_be_split_forall_dep (fun x -> cond1 x `prop_and` cond2 x) p r))
=
  Classical.forall_intro_3 (fun x y z -> Classical.move_requires (can_be_split_trans x y) z)

let can_be_split_forall_dep_congr_l
  (#a: Type)
  (cond: a -> prop)
  (p q r: post_t a)
: Lemma
  (requires (can_be_split_forall_dep cond p q))
  (ensures (can_be_split_forall_dep cond (fun x -> p x `star` r x) (fun x -> q x `star` r x)))
=
  Classical.forall_intro_3 (fun x y z -> Classical.move_requires (can_be_split_congr_l x y) z)

let can_be_split_forall_dep_congr_r
  (#a: Type)
  (cond: a -> prop)
  (p q r: post_t a)
: Lemma
  (requires (can_be_split_forall_dep cond p q))
  (ensures (can_be_split_forall_dep cond (fun x -> r x `star` p x) (fun x -> r x `star` q x)))
=
  Classical.forall_intro_3 (fun x y z -> Classical.move_requires (can_be_split_congr_r x y) z)

/// To simplify the implementation of the framing tactic, dependent equivalence
/// is defined as a double dependent implication
let equiv_forall (#a:Type) (t1 t2:post_t a) : Type0
  = t1 `can_be_split_forall` t2 /\ t2 `can_be_split_forall` t1

/// This equivalence models a context restriction at the end of a Steel computation;
/// note that t2 does not depend on the value of type `a`, but the two vprops must be
///  equivalent
let can_be_split_post (#a #b:Type) (t1:a -> post_t b) (t2:post_t b) =
  forall (x:a). equiv_forall (t1 x) t2

/// Lifting the equivalence relation to vprops. Two vprops are equivalent if the underlying slprops
/// are equivalent
val equiv (p q:vprop) : prop
/// Revealing the definition of vprop equivalence when needed for proof purposes.
/// In other cases, the predicate is abstract
val reveal_equiv (p q:vprop) : Lemma (p `equiv` q <==> hp_of p `Mem.equiv` hp_of q)

(* A restricted view of the heap,
   that only allows to access selectors of the current slprop *)

let rmem' (pre:vprop) =
  FExt.restricted_g_t
  (r0:vprop{can_be_split pre r0})
  (fun r0 -> normal (t_of r0))

/// Ensuring that rmems encapsulate the structure induced by the separation logic star
val valid_rmem (#frame:vprop) (h:rmem' frame) : prop

unfold
let rmem (pre:vprop) = h:rmem' pre{valid_rmem h}

/// Exposing the definition of mk_rmem to better normalize Steel VCs
unfold noextract
let unrestricted_mk_rmem (r:vprop) (h:hmem r) = fun (r0:vprop{r `can_be_split` r0}) ->
  can_be_split_interp r r0 h;
  sel_of r0 h

[@@ __inner_steel_reduce__]
noextract
let mk_rmem' (r:vprop) (h:hmem r) : Tot (rmem' r) =
   FExt.on_dom_g
     (r0:vprop{r `can_be_split` r0})
     (unrestricted_mk_rmem r h)

val lemma_valid_mk_rmem (r:vprop) (h:hmem r) : Lemma (valid_rmem (mk_rmem' r h))

[@@ __inner_steel_reduce__]
noextract
let mk_rmem (r:vprop) (h:hmem r) : Tot (rmem r) =
  lemma_valid_mk_rmem r h;
  mk_rmem' r h

val reveal_mk_rmem (r:vprop) (h:hmem r) (r0:vprop{r `can_be_split` r0})
  : Lemma (ensures reveal_can_be_split(); (mk_rmem r h) r0 == sel_of r0 h)

(* Logical pre and postconditions can only access the restricted view of the heap *)

type req_t (pre:pre_t) = rmem pre -> Type0
type ens_t (pre:pre_t) (a:Type) (post:post_t a) =
  rmem pre -> (x:a) -> rmem (post x) -> Type0

(* Empty assertion *)
val emp : vprop

/// When needed for proof purposes, the empty assertion is a direct lift of the
/// empty assertion from Steel.Memory
val reveal_emp (_:unit) : Lemma (hp_of emp == Mem.emp /\ t_of emp == unit)

/// Lifting pure predicates to vprop
[@@__steel_reduce__]
unfold let pure (p:prop) = to_vprop (pure p)

/// Framing predicates for the Steel effect. If the current computation has already
/// been framed, then the additional frame is the empty predicate
let maybe_emp (framed:bool) (frame:pre_t) = if framed then frame == emp else True
/// Dependent version of the above predicate, usable in dependent postconditions
let maybe_emp_dep (#a:Type) (framed:bool) (frame:post_t a) =
  if framed then (forall x. frame x == emp) else True

(* focus_rmem is an additional restriction of our view of memory.
   We expose it here to be able to reduce through normalization;
   Any valid application of focus_rmem h will be reduced to the application of h *)

[@@ __steel_reduce__]
unfold
let unrestricted_focus_rmem (#r:vprop) (h:rmem r) (r0:vprop{r `can_be_split` r0})
  = fun (r':vprop{can_be_split r0 r'}) -> can_be_split_trans r r0 r'; h r'

[@@ __inner_steel_reduce__]
let focus_rmem' (#r: vprop) (h: rmem r) (r0: vprop{r `can_be_split` r0}) : Tot (rmem' r0)
 = FExt.on_dom_g
   (r':vprop{can_be_split r0 r'})
   (unrestricted_focus_rmem h r0)

val lemma_valid_focus_rmem (#r:vprop) (h:rmem r) (r0:vprop{r `can_be_split` r0})
  : Lemma (valid_rmem (focus_rmem' h r0))

[@@ __inner_steel_reduce__]
let focus_rmem (#r:vprop) (h:rmem r) (r0:vprop{r `can_be_split` r0}) : Tot (rmem r0) =
  lemma_valid_focus_rmem h r0;
  focus_rmem' h r0

/// Exposing that calling focus_rmem on the current context corresponds to an equality
let focus_rmem_refl (r:vprop) (h:rmem r)
  : Lemma (focus_rmem #r h r == h)
  = FStar.FunctionalExtensionality.extensionality_g _ _ (focus_rmem #r h r) h

open FStar.Tactics.V2

/// State that all "atomic" subresources have the same selectors on both views.
/// The predicate has the __steel_reduce__ attribute, ensuring that VC normalization
/// will reduce it to a conjunction of equalities on atomic subresources
/// This predicate is also marked as `strict_on_arguments` on [frame], ensuring that
/// it will not be reduced when the frame is symbolic
/// Instead, the predicate will be rewritten to an equality using `lemma_frame_equalities` below
[@@ __steel_reduce__; strict_on_arguments [0]]
let rec frame_equalities'
  (frame:vprop)
  (h0:rmem frame) (h1:rmem frame) : Type0
  = begin match frame with
    | VUnit p -> h0 frame == h1 frame
    | VStar p1 p2 ->
        can_be_split_star_l p1 p2;
        can_be_split_star_r p1 p2;

        let h01 = focus_rmem h0 p1 in
        let h11 = focus_rmem h1 p1 in

        let h02 = focus_rmem h0 p2 in
        let h12 = focus_rmem h1 p2 in


        frame_equalities' p1 h01 h11 /\
        frame_equalities' p2 h02 h12
    end

/// This lemma states that frame_equalities is the same as an equality on the top-level frame.
/// The uncommon formulation with an extra [p] is needed to use in `rewrite_with_tactic`,
/// where the goal is of the shape `frame_equalities frame h0 h1 == ?u`
/// The rewriting happens below, in `frame_vc_norm`
val lemma_frame_equalities (frame:vprop) (h0:rmem frame) (h1:rmem frame) (p:Type0)
  : Lemma
    (requires (h0 frame == h1 frame) == p)
    (ensures frame_equalities' frame h0 h1 == p)

/// A special case for frames about emp.
val lemma_frame_emp (h0:rmem emp) (h1:rmem emp) (p:Type0)
  : Lemma (requires True == p)
          (ensures frame_equalities' emp h0 h1 == p)

/// A variant of conjunction elimination, suitable to the equality goals during rewriting
val elim_conjunction (p1 p1' p2 p2':Type0)
  : Lemma (requires p1 == p1' /\ p2 == p2')
          (ensures (p1 /\ p2) == (p1' /\ p2'))

/// Normalization and rewriting step for generating frame equalities.
/// The frame_equalities function has the strict_on_arguments attribute on the [frame],
/// ensuring that it is not reduced when the frame is symbolic.
/// When that happens, we want to replace frame_equalities by an equality on the frame,
/// mimicking reduction
[@@plugin]
let frame_vc_norm () : Tac unit = with_compat_pre_core 0 (fun _ ->
  // Do not normalize mk_rmem/focus_rmem to simplify application of
  // the reflexivity lemma on frame_equalities'
  norm [delta_attr [`%__steel_reduce__];
    delta_only [`%Mkvprop'?.t; `%Mkvprop'?.hp; `%Mkvprop'?.sel;
      `%FStar.Algebra.CommMonoid.Equiv.__proj__CM__item__mult;
      `%FStar.Algebra.CommMonoid.Equiv.__proj__CM__item__unit];
    delta_qualifier ["unfold"];
    iota;zeta;primops; simplify];

  // After reduction, the term to rewrite might be of the shape
  // (frame_equalities' ... /\ frame_equalities' .. /\ ...) == ?u,
  // with some frame_equalities' possibly already fully reduced
  // We repeatedly split the clause and extract the term on the left
  // to generate equalities on atomic subresources
  ignore (repeat (fun _ ->
    // Try to split the conjunction. If there is no conjunction, we exit the repeat
    apply_lemma (`elim_conjunction);
    // Dismiss the two uvars created for the RHS, they'll be solved by unification
    dismiss ();
    dismiss ();
    // The first goal is the left conjunction
    split ();
    // Removes the frame equality if it is about emp
    or_else (fun _ -> apply_lemma (`lemma_frame_emp); dismiss()) (fun _ -> ());
    // Rewrites the frame_equalities if it wasn't yet reduced
    or_else (fun _ -> apply_lemma (`lemma_frame_equalities); dismiss ()) (fun _ -> ());
    norm normal_steps;
    // Finally solve the uvar, finishing the rewriting for this clause
    trefl ()
  ));

  // Removes the frame equality if it is about emp
  or_else (fun _ -> apply_lemma (`lemma_frame_emp); dismiss()) (fun _ -> ());

  // We do not have conjunctions anymore, we try to apply the frame_equalities rewriting
  // If it fails, the frame was not symbolic, so there is nothing to do
  or_else (fun _ -> apply_lemma (`lemma_frame_equalities); dismiss ()) (fun _ -> ());
  norm normal_steps;
  trefl ())

[@@ __steel_reduce__]
unfold
let frame_equalities
  (frame:vprop)
  (h0:rmem frame) (h1:rmem frame) : prop
  = rewrite_with_tactic frame_vc_norm (frame_equalities' frame h0 h1)

/// More lemmas about the abstract can_be_split predicates, to be used as
/// rewriting rules in the tactic below
val can_be_split_dep_refl (p:vprop)
: Lemma (can_be_split_dep true_p p p)

val equiv_can_be_split (p1 p2:vprop) : Lemma
  (requires p1 `equiv` p2)
  (ensures p1 `can_be_split` p2)

val intro_can_be_split_frame (p q:vprop) (frame:vprop)
  : Lemma (requires q `equiv` (p `star` frame))
          (ensures can_be_split q p /\ True)

val can_be_split_post_elim (#a #b:Type) (t1:a -> post_t b) (t2:post_t b)
  : Lemma (requires (forall (x:a) (y:b). t1 x y `equiv` t2 y))
          (ensures t1 `can_be_split_post` t2)

val equiv_forall_refl (#a:Type) (t:post_t a)
  : Lemma (t `equiv_forall` t)

val equiv_forall_elim (#a:Type) (t1 t2:post_t a)
  : Lemma (requires (forall (x:a). t1 x `equiv` t2 x))
          (ensures t1 `equiv_forall` t2)

open FStar.Tactics.CanonCommMonoidSimple.Equiv

(* equiv is an equivalence relation on vprops *)

/// Lemmas establishing the equivalence properties on equiv
val equiv_refl (x:vprop) : Lemma (equiv x x)

val equiv_sym (x y:vprop) : Lemma
  (requires equiv x y)
  (ensures equiv y x)

val equiv_trans (x y z:vprop) : Lemma
  (requires equiv x y /\ equiv y z)
  (ensures equiv x z)

module CE = FStar.Algebra.CommMonoid.Equiv

/// Equiv is an equivalence relation for vprops elements
inline_for_extraction noextract let req : CE.equiv vprop =
  CE.EQ equiv
     equiv_refl
     equiv_sym
     equiv_trans

(* Star induces a commutative monoid for the equiv equivalence relation *)

/// Lemmas establishing the commutative monoid properties
val cm_identity (x:vprop) : Lemma ((emp `star` x) `equiv` x)

val star_commutative (p1 p2:vprop)
  : Lemma ((p1 `star` p2) `equiv` (p2 `star` p1))

val star_associative (p1 p2 p3:vprop)
  : Lemma (((p1 `star` p2) `star` p3)
           `equiv`
           (p1 `star` (p2 `star` p3)))

val star_congruence (p1 p2 p3 p4:vprop)
  : Lemma (requires p1 `equiv` p3 /\ p2 `equiv` p4)
          (ensures (p1 `star` p2) `equiv` (p3 `star` p4))

/// Star induces a commutative monoid on vprops
[@__steel_reduce__]
inline_for_extraction noextract let rm : CE.cm vprop req =
  CE.CM emp
     star
     cm_identity
     star_associative
     star_commutative
     star_congruence

(*** Vprop combinators ***)

(* Refining a vprop with a selector predicate *)

/// Separation logic predicate stating the validity of a vprop with an additional refinement on its selector
val vrefine_hp (v: vprop) (p: (normal (t_of v) -> Tot prop)) : Tot (slprop u#1)

/// Exposing the validity of the above predicate when needed for proof purposes
val interp_vrefine_hp (v: vprop) (p: (normal (t_of v) -> Tot prop)) (m: mem) : Lemma
  (interp (vrefine_hp v p) m <==> (interp (hp_of v) m /\ p (sel_of v m)))

/// Selector type for a refined vprop
[@__steel_reduce__]
let vrefine_t (v: vprop) (p: (normal (t_of v) -> Tot prop)) : Tot Type
= (x: t_of v {p x})

/// Selector of a refined vprop. Returns a value which satisfies the refinement predicate
val vrefine_sel (v: vprop) (p: (normal (t_of v) -> Tot prop)) : Tot (selector (vrefine_t v p) (vrefine_hp v p))

/// Exposing the definition of the refined selector
val vrefine_sel_eq (v: vprop) (p: (normal (t_of v) -> Tot prop)) (m: Mem.hmem (vrefine_hp v p)) : Lemma
  (
    interp (hp_of v) m /\
    vrefine_sel v p m == sel_of v m
  )
//  [SMTPat ((vrefine_sel v p) m)] // FIXME: this pattern causes Z3 "wrong number of argument" errors

/// Combining the above pieces to define a vprop refined by a selector prediacte
[@__steel_reduce__]
let vrefine' (v: vprop) (p: (normal (t_of v) -> Tot prop)) : Tot vprop' = {
  hp = vrefine_hp v p;
  t = vrefine_t v p;
  sel = vrefine_sel v p;
}

[@__steel_reduce__]
let vrefine (v: vprop) (p: (normal (t_of v) -> Tot prop)) = VUnit (vrefine' v p)

(* Dependent star for vprops *)

/// Separation logic predicate corresponding to a dependent star,
/// where the second predicate depends on the selector value of the first
val vdep_hp (v: vprop) (p: ( (t_of v) -> Tot vprop)) : Tot (slprop u#1)

/// Exposing the validity of the above predicate when needed for proof purposes
val interp_vdep_hp (v: vprop) (p: ( (t_of v) -> Tot vprop)) (m: mem) : Lemma
  (interp (vdep_hp v p) m <==> (interp (hp_of v) m /\ interp (hp_of v `Mem.star` hp_of (p (sel_of v m))) m))

/// Helper to define the selector type of the second component of the dependent star
let vdep_payload
  (v: vprop) (p: ( (t_of v) -> Tot vprop))
  (x: t_of v)
: Tot Type
= t_of (p x)

/// Selector type for the dependent star: A dependent tuple, where the second component's type depends on the first vprop
let vdep_t (v: vprop) (p: ( (t_of v) -> Tot vprop)) : Tot Type
= dtuple2 (t_of v) (vdep_payload v p)

/// Selector for the dependent star
val vdep_sel (v: vprop) (p: ( (t_of v) -> Tot vprop)) : Tot (selector (vdep_t v p) (vdep_hp v p))

/// Exposing the definition of the dependent star's selector when needed for proof purposes
val vdep_sel_eq (v: vprop) (p: ( (t_of v) -> Tot vprop)) (m: Mem.hmem (vdep_hp v p)) : Lemma
  (
    interp (hp_of v) m /\
    begin let x = sel_of v m in
      interp (hp_of (p x)) m /\
      vdep_sel v p m == (| x, sel_of (p x) m |)
    end
  )

/// Combining the elements above to create a dependent star vprop
[@__steel_reduce__]
let vdep' (v: vprop) (p: ( (t_of v) -> Tot vprop)) : Tot vprop' = {
  hp = vdep_hp v p;
  t = vdep_t v p;
  sel = vdep_sel v p;
}

[@__steel_reduce__]
let vdep (v: vprop) (p: ( (t_of v) -> Tot vprop)) = VUnit (vdep' v p)

(* Selector rewrite combinator *)

/// The selector of a rewrite combinator applies a function `f` to the current selector of a vprop.
val vrewrite_sel (v: vprop) (#t: Type) (f: (normal (t_of v) -> GTot t)) : Tot (selector t (normal (hp_of v)))

/// Exposing the definition of the above selector
val vrewrite_sel_eq (v: vprop) (#t: Type) (f: (normal (t_of v) -> GTot t)) (h: Mem.hmem (normal (hp_of v))) : Lemma
  ((vrewrite_sel v f <: selector' _ _) h == f ((normal (sel_of v) <: selector' _ _) h))
//  [SMTPat (vrewrite_sel v f h)] // FIXME: this pattern causes Z3 "wrong number of argument" errors


/// Combining the above elements to create a rewrite vprop
[@__steel_reduce__]
let vrewrite' (v: vprop) (#t: Type) (f: (normal (t_of v) -> GTot t)) : Tot vprop' = {
  hp = normal (hp_of v);
  t = t;
  sel = vrewrite_sel v f;
}

[@__steel_reduce__]
let vrewrite (v: vprop) (#t: Type) (f: (normal (t_of v) -> GTot t)) : Tot vprop = VUnit (vrewrite' v f)

(*** Framing tactic ***)

(* Specialize visit_tm from the standard F* tactic library to reimplement name_appears_in.
   AF: As of Jan 14, 2021, calling name_appears_in from FStar.Tactics.Derived leads to a segfault *)

exception Appears

let on_sort_binder (f : term -> Tac unit) (b:binder) : Tac unit =
  f b.sort

let rec visit_tm (ff : term -> Tac unit) (t : term) : Tac unit =
  let tv = inspect t in
  (match tv with
  | Tv_FVar _
  | Tv_UInst _ _
  | Tv_Var _
  | Tv_BVar _ -> ()

  | Tv_Type _ -> ()
  | Tv_Const c -> ()
  | Tv_Uvar i u -> ()
  | Tv_Unsupp  -> ()
  | Tv_Unknown -> ()
  | Tv_Arrow b c ->
      on_sort_binder ff b;
      visit_comp ff c
  | Tv_Abs b t ->
      let b = on_sort_binder (visit_tm ff) b in
      visit_tm ff t

  | Tv_App l (r, q) ->
       visit_tm ff l;
       visit_tm ff r

  | Tv_Refine b r ->
      on_sort_binder ff b;
      visit_tm ff r

  | Tv_Let r attrs b def t ->
      on_sort_binder ff b;
      visit_tm ff def;
      visit_tm ff t

  | Tv_Match sc _ brs ->
      visit_tm ff sc;
      iter (visit_br ff) brs

  | Tv_AscribedT e t topt _ ->
      visit_tm ff e;
      visit_tm ff t

  | Tv_AscribedC e c topt _ ->
      visit_tm ff e

  ); ff t

and visit_br (ff : term -> Tac unit) (b:branch) : Tac unit =
  let (p, t) = b in
  visit_tm ff t

and visit_comp (ff : term -> Tac unit) (c : comp) : Tac unit =
  let cv = inspect_comp c in
  match cv with
  | C_Total ret -> visit_tm ff ret
  | C_GTotal ret -> visit_tm ff ret

  | C_Lemma pre post pats ->
      visit_tm ff pre;
      visit_tm ff post;
      visit_tm ff pats

  | C_Eff us eff res args decrs ->
      visit_tm ff res;
      iter (fun (a, q) -> visit_tm ff a) args;
      iter (visit_tm ff) decrs

/// Decides whether a top-level name [nm] syntactically
/// appears in the term [t].
let name_appears_in (nm:name) (t:term) : Tac bool =
  let ff (t : term) : Tac unit =
    match inspect t with
    | Tv_FVar fv -> if inspect_fv fv = nm then raise Appears
    | t -> ()
  in
  try ignore (visit_tm ff t); false with
  | Appears -> true
  | e -> raise e

/// Checks whether term [t] appears in term [i]
let term_appears_in (t:term) (i:term) : Tac bool =
  name_appears_in (explode_qn (term_to_string t)) i


/// We define a small language to handle arbitrary separation logic predicates.
/// Separation logic predicates are encoded as atoms for which equality is decidable,
/// here represented as integers
let atom : eqtype = int

let rec atoms_to_string (l:list atom) = match l with
  | [] -> ""
  | hd::tl -> string_of_int hd ^ " " ^ atoms_to_string tl

/// Reflecting the structure of our separation logic on atmos
type exp : Type =
  | Unit : exp
  | Mult : exp -> exp -> exp
  | Atom : atom -> exp

/// A map from atoms to the terms they represent.
/// The second component of the term corresponds to a default element,
/// ensuring we never raise an exception when trying to access an element in the map
let amap (a:Type) = list (atom * a) * a

/// An empty atom map: The list map is empty
let const (#a:Type) (xa:a) : amap a = ([], xa)

/// Accessing an element in the atom map

// We reimplement List.Tot.Base.assoc because we need our tactic to normalize it,
// but we don't want to normalize user-provided instances
let rec my_assoc (#key: eqtype) (#value: Type) (k: key) (dict: list (key & value)) : Pure (option value)
  (requires True)
  (ensures (fun res -> res == List.Tot.assoc k dict))
= match dict with
  | [] -> None
  | (k', v') :: q -> if k = k' then Some v' else my_assoc k q


let select (#a:Type) (x:atom) (am:amap a) : Tot a =
  match my_assoc #atom #a x (fst am) with
  | Some a -> a
  | _ -> snd am

/// Updating the atom map. Since select finds the first element corresponding to
/// the atom in the list and we do not have any remove function,
/// we can simply append the new element at the head without removing any possible
/// previous element
let update (#a:Type) (x:atom) (xa:a) (am:amap a) : amap a =
  (x, xa)::fst am, snd am

/// Check whether the current term is an unresolved vprop unification variable.
/// This can happen if either it is a uvar, or it is an unresolved dependent
/// vprop uvar which is applied to some argument
let is_uvar (t:term) : Tac bool = match inspect t with
  | Tv_Uvar _ _ -> true
  | Tv_App _ _ ->
      let hd, args = collect_app t in
      Tv_Uvar? (inspect hd)
  | _ -> false

/// For a given term t, collect all terms in the list l with the same head symbol
let rec get_candidates (t:term) (l:list term) : Tac (list term) =
  let name, _ = collect_app t in
  match l with
  | [] -> []
  | hd::tl ->
      let n, _ = collect_app hd in
      if term_eq_old n name then (
        hd::(get_candidates t tl)
      ) else get_candidates t tl

/// Try to remove a term that is exactly matching, not just that can be unified
let rec trivial_cancel (t:atom) (l:list atom) =
  match l with
  | [] -> false, l
  | hd::tl ->
      if hd = t then
        // These elements match, we remove them
        true, tl
      else (let b, res = trivial_cancel t tl in b, hd::res)

/// Call trivial_cancel on all elements of l1.
/// The first two lists returned are the remainders of l1 and l2.
/// The last two lists are the removed parts of l1 and l2, with
/// the additional invariant that they are equal
let rec trivial_cancels (l1 l2:list atom) (am:amap term)
  : Tac (list atom * list atom * list atom * list atom) =
  match l1 with
  | [] -> [], l2, [], []
  | hd::tl ->
      let b, l2'   = trivial_cancel hd l2 in
      let l1', l2', l1_del, l2_del = trivial_cancels tl l2' am in
      (if b then l1' else hd::l1'), l2',
      (if b then hd::l1_del else l1_del), (if b then hd::l2_del else l2_del)

exception Failed
exception Success

/// Helper to print the terms corresponding to the current list of atoms
let rec print_atoms (l:list atom) (am:amap term) : Tac string =
  match l with
  | [] -> ""
  | [hd] -> term_to_string (select hd am)
  | hd::tl -> term_to_string (select hd am) ^ " * " ^ print_atoms tl am

/// For a list of candidates l, count the number that can unify with t.
/// Does not try to unify with a uvar, this will be done at the very end.
/// Tries to unify with slprops with a different head symbol, it might
/// be an abbreviation
let rec try_candidates (t:atom) (l:list atom) (am:amap term) : Tac (atom * int) =
  match l with
  | [] -> t, 0
  | hd::tl ->
      if is_uvar (select hd am) then (try_candidates t tl am)
      else
        // Encapsulate unify in a try/with to ensure unification is not actually performed
        let res = try if unify (select t am) (select hd am) then raise Success else raise Failed
                  with | Success -> true | _ -> false in
        let t', n' = try_candidates t tl am in
        if res && hd <> t' then hd, 1 + n'  else t', n'

/// Remove the given term from the list. Only to be called when
/// try_candidates succeeded
let rec remove_from_list (t:atom) (l:list atom) : Tac (list atom) =
  match l with
  | [] -> fail "atom in remove_from_list not found: should not happen"; []
  | hd::tl -> if t = hd then tl else hd::remove_from_list t tl

/// Check if two lists of slprops are equivalent by recursively calling
/// try_candidates.
/// Assumes that only l2 contains terms with the head symbol unresolved.
/// It returns all elements that were not resolved during this iteration *)
let rec equivalent_lists_once (l1 l2 l1_del l2_del:list atom) (am:amap term)
  : Tac (list atom * list atom * list atom * list atom) =
  match l1 with
  | [] -> [], l2, l1_del, l2_del
  | hd::tl ->
    let t, n = try_candidates hd l2 am in
    if n = 1 then (
      let l2 = remove_from_list t l2 in
      equivalent_lists_once tl l2 (hd::l1_del) (t::l2_del) am
    ) else (
      // Either too many candidates for this scrutinee, or no candidate but the uvar
      let rem1, rem2, l1'_del, l2'_del = equivalent_lists_once tl l2 l1_del l2_del am in
      hd::rem1, rem2, l1'_del, l2'_del
    )

/// Check if two lists of slprops are equivalent by recursively calling
/// try_candidates by iterating on l2.
/// Assumes that only l2 contains terms with the head symbol unresolved.
/// It returns all elements that were not resolved during this iteration *)
/// This is very close to equivalent_lists_once above, but helps making progress
/// when l1 contains syntactically equal candidates
let rec equivalent_lists_once_l2 (l1 l2 l1_del l2_del:list atom) (am:amap term)
  : Tac (list atom * list atom * list atom * list atom) =
  match l2 with
  | [] -> l1, [], l1_del, l2_del
  | hd::tl ->
    if is_uvar (select hd am) then
      // We do not try to match the vprop uvar
      let rem1, rem2, l1'_del, l2'_del = equivalent_lists_once_l2 l1 tl l1_del l2_del am in
      rem1, hd::rem2, l1'_del, l2'_del
    else (
      let t, n = try_candidates hd l1 am in
      if n = 1 then (
        let l1 = remove_from_list t l1 in
        equivalent_lists_once_l2 l1 tl (t::l1_del) (hd::l2_del) am
      ) else (
        // Either too many candidates for this scrutinee, or no candidate but the uvar
        let rem1, rem2, l1'_del, l2'_del = equivalent_lists_once_l2 l1 tl l1_del l2_del am in
        rem1, hd::rem2, l1'_del, l2'_del
      )
    )


let get_head (l:list atom) (am:amap term) : term = match l with
  | [] -> `()
  | hd::_ -> select hd am

/// Checks whether the list of atoms [l] only contains one unresolved uvar
let is_only_uvar (l:list atom) (am:amap term) : Tac bool =
  if List.Tot.Base.length l = 1 then is_uvar (select (List.Tot.Base.hd l) am)
  else false

/// Assumes that u is a uvar, checks that all variables in l can be unified with it.
/// Later in the tactic, the uvar will be unified to a star of l *)
let rec try_unifying_remaining (l:list atom) (u:term) (am:amap term) : Tac unit =
  match l with
  | [] -> ()
  | hd::tl ->
      try if unify u (select hd am) then raise Success else raise Failed with
      | Success -> try_unifying_remaining tl u am
      | _ -> fail ("could not find candidate for scrutinee " ^ term_to_string (select hd am))

/// Is SMT rewriting enabled for this binder
let is_smt_binder (b:binder) : Tac bool =
  let l = b.attrs in
  not (List.Tot.isEmpty (filter (fun t -> is_fvar t (`%smt_fallback)) l))

/// Creates a new term, where all arguments where SMT rewriting is enabled have been replaced
/// by fresh, unconstrained unification variables
let rec new_args_for_smt_attrs (env:env) (l:list argv) (ty:typ) : Tac (list argv * list term) =
  let fresh_ghost_uvar ty =
    let e = cur_env () in
    ghost_uvar_env e ty
  in
  match l, inspect_unascribe ty with
  | (arg, aqualv)::tl, Tv_Arrow binder comp ->
    let needs_smt = is_smt_binder binder in
    let new_hd =
      if needs_smt then (
        let arg_ty = tc env arg in
        let uvar = fresh_ghost_uvar arg_ty in
        unshelve uvar;
        flip ();
        (uvar, aqualv)
      ) else (arg, aqualv)
    in
    begin
    let ty2 =
      match inspect_comp comp with
      | C_Total ty2 -> ty2
      | C_Eff _ eff_name ty2 _ _ ->
        if eff_name = ["Prims"; "Tot"]
        then ty2
        else fail "computation type not supported in definition of slprops"
      | _ -> fail "computation type not supported in definition of slprops" in

    let tl_argv, tl_terms = new_args_for_smt_attrs env tl ty2 in
    new_hd::tl_argv, (if needs_smt then arg::tl_terms else tl_terms)
    end
  | [], Tv_FVar fv -> [], []
  | _ -> fail "should not happen. Is an slprop partially applied?"

/// Rewrites all terms in the context to enable SMT rewriting through the use of fresh, unconstrained unification variables
let rewrite_term_for_smt (env:env) (am:amap term * list term) (a:atom) : Tac (amap term * list term)
  = let am, prev_uvar_terms = am in
    let term = select a am in
    let hd, args = collect_app term in
    let t = tc env hd in
    let new_args, uvar_terms = new_args_for_smt_attrs env args t in
    let new_term = mk_app hd new_args in
    update a new_term am, List.Tot.append uvar_terms prev_uvar_terms

/// User-facing error message when the framing tactic fails
let fail_atoms (#a:Type) (l1 l2:list atom) (am:amap term) : Tac a
  = fail ("could not find a solution for unifying\n" ^ print_atoms l1 am ^ "\nand\n" ^ print_atoms l2 am)

/// Variant of equivalent_lists' below to be called once terms have been rewritten to allow SMT rewriting.
/// If unification succeeds and we have unicity of the solution, this tactic will succeed,
/// and ultimately create an SMT guard that the two terms are actually equal
let rec equivalent_lists_fallback (n:nat) (l1 l2 l1_del l2_del:list atom) (am:amap term)
  : Tac (list atom * list atom * bool) =
  match l1 with
  | [] -> begin match l2 with
    | [] -> (l1_del, l2_del, false)
    | [hd] ->
      // Succeed if there is only one uvar left in l2, which can be therefore
      // be unified with emp
      if is_uvar (select hd am) then (
        // xsdenote is left associative: We put hd at the top to get
        // ?u `star` p <==> emp `star` p
        (l1_del, hd :: l2_del, true))
      else fail ("could not find candidates for " ^ term_to_string (get_head l2 am))
    | _ -> fail ("could not find candidates for " ^ term_to_string (get_head l2 am))
    end
  | _ ->
    if is_only_uvar l2 am then (
      // Terms left in l1, but only a uvar left in l2.
      // Put all terms left at the end of l1_rem, so that they can be unified
      // with exactly the uvar because of the structure of xsdenote
      try_unifying_remaining l1 (get_head l2 am) am;
      l1_del `List.Tot.append` l1, l2_del `List.Tot.append` l2, false
    ) else
      let rem1, rem2, l1_del', l2_del' = equivalent_lists_once l1 l2 l1_del l2_del am in
      let n' = List.Tot.length rem1 in
      if n' >= n then
        // Should always be smaller or equal to n
        // If it is equal, no progress was made.
        fail_atoms rem1 rem2 am
      else equivalent_lists_fallback n' rem1 rem2 l1_del' l2_del' am

/// Iterates over all terms in [l2] to prepare them for unification with SMT rewriting
let replace_smt_uvars (l1 l2:list atom) (am:amap term) : Tac (amap term * list term)
  = let env = cur_env () in
    fold_left (rewrite_term_for_smt env) (am, []) l2

/// Recursively calls equivalent_lists_once.
/// Stops when we're done with unification, or when we didn't make any progress
/// If we didn't make any progress, we have too many candidates for some terms.
/// Accumulates rewritings of l1 and l2 in l1_del and l2_del, with the invariant
/// that the two lists are unifiable at any point
/// The boolean indicates if there is a leftover empty frame
let rec equivalent_lists' (n:nat) (use_smt:bool) (l1 l2 l1_del l2_del:list atom) (am:amap term)
  : Tac (list atom * list atom * bool * list term) =
  match l1 with
  | [] -> begin match l2 with
    | [] -> (l1_del, l2_del, false, [])
    | [hd] ->
      // Succeed if there is only one uvar left in l2, which can be therefore
      // be unified with emp
      if is_uvar (select hd am) then (
        // xsdenote is left associative: We put hd at the top to get
        // ?u `star` p <==> emp `star` p
        (l1_del, hd :: l2_del, true, []))
      else fail ("could not find candidates for " ^ term_to_string (get_head l2 am))
    | _ -> fail ("could not find candidates for " ^ term_to_string (get_head l2 am))
    end
  | _ ->
    if is_only_uvar l2 am then (
      // Terms left in l1, but only a uvar left in l2.
      // Put all terms left at the end of l1_rem, so that they can be unified
      // with exactly the uvar because of the structure of xsdenote
      try_unifying_remaining l1 (get_head l2 am) am;
      l1_del `List.Tot.append` l1, l2_del `List.Tot.append` l2, false, []
    ) else
      let rem1, rem2, l1_del', l2_del' = equivalent_lists_once l1 l2 l1_del l2_del am in
      let n' = List.Tot.length rem1 in
      if n' >= n then (
        // Try to make progress by matching non-uvars of l2 with candidates in l1
        let rem1, rem2, l1_del', l2_del' = equivalent_lists_once_l2 rem1 rem2 l1_del' l2_del' am in
        let n' = List.Tot.length rem1 in
        if n' >= n then (
          // Should always be smaller or equal to n
          // If it is equal, no progress was made.
          if use_smt then
            // SMT fallback is allowed
            let new_am, uvar_terms  = replace_smt_uvars rem1 rem2 am in
            let l1_f, l2_f, b = equivalent_lists_fallback n' rem1 rem2 l1_del' l2_del' new_am in
            l1_f, l2_f, b, uvar_terms
          else fail_atoms rem1 rem2 am
        ) else equivalent_lists' n' use_smt rem1 rem2 l1_del' l2_del' am
      ) else equivalent_lists' n' use_smt rem1 rem2 l1_del' l2_del' am

/// Checks if term for atom t unifies with fall uvars in l
let rec unifies_with_all_uvars (t:term) (l:list atom) (am:amap term) : Tac bool =
  match l with
  | [] -> true
  | hd::tl ->
      if unifies_with_all_uvars t tl am then (
        // Unified with tail, try this term
        let hd_t = select hd am in
        if is_uvar hd_t then (
          // The head term is a uvar, try unifying
          try if unify t hd_t then raise Success else raise Failed
          with | Success -> true | _ -> false
        ) else true // The uvar is not a head term, we do not need to try it
      ) else false

/// Puts all terms in l1 that cannot unify with the uvars in l2 at the top:
/// They need to be solved first
let rec most_restricted_at_top (l1 l2:list atom) (am:amap term) : Tac (list atom) =
  match l1 with
  | [] -> []
  | hd::tl ->
    if unifies_with_all_uvars (select hd am) l2 am then (most_restricted_at_top tl l2 am) `List.Tot.append` [hd]
    else hd::(most_restricted_at_top tl l2 am)

/// Core AC-unification tactic.
/// First remove all trivially equal terms, then try to decide equivalence.
/// Assumes that l1 does not contain any vprop uvar.
/// If it succeeds, returns permutations of l1, l2, and a boolean indicating
/// if l2 has a trailing empty frame to be unified
let equivalent_lists (use_smt:bool) (l1 l2:list atom) (am:amap term)
  : Tac (list atom * list atom * bool * list term)
= let l1, l2, l1_del, l2_del = trivial_cancels l1 l2 am in
  let l1 = most_restricted_at_top l1 l2 am in
  let n = List.Tot.length l1 in
  let l1_del, l2_del, emp_frame, uvar_terms = equivalent_lists' n use_smt l1 l2 l1_del l2_del am in
  l1_del, l2_del, emp_frame, uvar_terms

(* Helpers to relate the actual terms to their representation as a list of atoms *)

open FStar.Reflection.V2.Derived.Lemmas

let rec list_to_string (l:list term) : Tac string =
  match l with
  | [] -> "end"
  | hd::tl -> term_to_string hd ^ " " ^ list_to_string tl

let rec mdenote_gen (#a:Type u#aa) (unit:a) (mult:a -> a -> a) (am:amap a) (e:exp) : a =
  match e with
  | Unit -> unit
  | Atom x -> select x am
  | Mult e1 e2 -> mult (mdenote_gen unit mult am e1) (mdenote_gen unit mult am e2)

let rec xsdenote_gen (#a:Type) (unit:a) (mult:a -> a -> a) (am:amap a) (xs:list atom) : a =
  match xs with
  | [] -> unit
  | [x] -> select x am
  | x::xs' -> mult (select x am) (xsdenote_gen unit mult am xs')

unfold
let mdenote (#a:Type u#aa) (eq:CE.equiv a) (m:CE.cm a eq) (am:amap a) (e:exp) : a =
  let open FStar.Algebra.CommMonoid.Equiv in
  mdenote_gen (CM?.unit m) (CM?.mult m) am e

unfold
let xsdenote (#a:Type) (eq:CE.equiv a) (m:CE.cm a eq) (am:amap a) (xs:list atom) : a =
  let open FStar.Algebra.CommMonoid.Equiv in
  xsdenote_gen (CM?.unit m) (CM?.mult m) am xs

// We reimplement List.Tot.Base.append because we need our tactic to normalize it,
// but we don't want to normalize user-provided instances

let rec my_append (#t: Type) (l1 l2: list t) : Pure (list t)
  (requires True)
  (ensures (fun res -> res == l1 `List.Tot.append` l2))
  (decreases l1)
= match l1 with
  | [] -> l2
  | a :: q -> a :: my_append q l2

let rec flatten (e:exp) : list atom =
  match e with
  | Unit -> []
  | Atom x -> [x]
  | Mult e1 e2 -> flatten e1 `my_append` flatten e2

let rec flatten_correct_aux (#a:Type) (eq:CE.equiv a) (m:CE.cm a eq) (am:amap a) (xs1 xs2:list atom)
  : Lemma (xsdenote eq m am (xs1 `my_append` xs2) `CE.EQ?.eq eq` CE.CM?.mult m (xsdenote eq m am xs1)
                                                                     (xsdenote eq m am xs2)) =
  let open FStar.Algebra.CommMonoid.Equiv in
  match xs1 with
  | [] ->
      CM?.identity m (xsdenote eq m am xs2);
      EQ?.symmetry eq (CM?.mult m (CM?.unit m) (xsdenote eq m am xs2)) (xsdenote eq m am xs2)
  | [x] -> (
      if (Nil? xs2)
      then (right_identity eq m (select x am);
            EQ?.symmetry eq (CM?.mult m (select x am) (CM?.unit m)) (select x am))
      else EQ?.reflexivity eq (CM?.mult m (xsdenote eq m am [x]) (xsdenote eq m am xs2)))
  | x::xs1' ->
      flatten_correct_aux eq m am xs1' xs2;
      EQ?.reflexivity eq (select x am);
      CM?.congruence m (select x am) (xsdenote eq m am (xs1' `my_append` xs2))
                       (select x am) (CM?.mult m (xsdenote eq m am xs1') (xsdenote eq m am xs2));
      CM?.associativity m (select x am) (xsdenote eq m am xs1') (xsdenote eq m am xs2);
      EQ?.symmetry eq (CM?.mult m (CM?.mult m (select x am) (xsdenote eq m am xs1')) (xsdenote eq m am xs2))
                      (CM?.mult m (select x am) (CM?.mult m (xsdenote eq m am xs1') (xsdenote eq m am xs2)));
      EQ?.transitivity eq (CM?.mult m (select x am) (xsdenote eq m am (xs1' `my_append` xs2)))
                          (CM?.mult m (select x am) (CM?.mult m (xsdenote eq m am xs1') (xsdenote eq m am xs2)))
                          (CM?.mult m (CM?.mult m (select x am) (xsdenote eq m am xs1')) (xsdenote eq m am xs2))

let rec flatten_correct (#a:Type) (eq:CE.equiv a) (m:CE.cm a eq) (am:amap a) (e:exp)
  : Lemma (mdenote eq m am e `CE.EQ?.eq eq` xsdenote eq m am (flatten e)) =
  let open FStar.Algebra.CommMonoid.Equiv in
  match e with
  | Unit -> EQ?.reflexivity eq (CM?.unit m)
  | Atom x -> EQ?.reflexivity eq (select x am)
  | Mult e1 e2 ->
      flatten_correct_aux eq m am (flatten e1) (flatten e2);
      EQ?.symmetry eq (xsdenote eq m am (flatten e1 `my_append` flatten e2))
                      (CM?.mult m (xsdenote eq m am (flatten e1)) (xsdenote eq m am (flatten e2)));
      flatten_correct eq m am e1;
      flatten_correct eq m am e2;
      CM?.congruence m (mdenote eq m am e1) (mdenote eq m am e2)
                       (xsdenote eq m am (flatten e1)) (xsdenote eq m am (flatten e2));
      EQ?.transitivity eq (CM?.mult m (mdenote eq m am e1) (mdenote eq m am e2))
                          (CM?.mult m (xsdenote eq m am (flatten e1)) (xsdenote eq m am (flatten e2)))
                          (xsdenote eq m am (flatten e1 `my_append` flatten e2))

let monoid_reflect (#a:Type) (eq:CE.equiv a) (m:CE.cm a eq) (am:amap a) (e1 e2:exp)
    (_ : squash (xsdenote eq m am (flatten e1) `CE.EQ?.eq eq` xsdenote eq m am (flatten e2)))
       : squash (mdenote eq m am e1 `CE.EQ?.eq eq` mdenote eq m am e2) =
    flatten_correct eq m am e1;
    flatten_correct eq m am e2;
    CE.EQ?.symmetry eq (mdenote eq m am e2) (xsdenote eq m am (flatten e2));
    CE.EQ?.transitivity eq
      (xsdenote eq m am (flatten e1))
      (xsdenote eq m am (flatten e2))
      (mdenote eq m am e2);
    CE.EQ?.transitivity eq
      (mdenote eq m am e1)
      (xsdenote eq m am (flatten e1))
      (mdenote eq m am e2)

// Here we sort the variable numbers

// We reimplement List.Tot.Base.sortWith because we need our tactic to normalize it,
// but we don't want to normalize user-provided instances
let rec my_partition (#a: Type) (f: (a -> Tot bool)) (l: list a)
  : Pure (list a & list a)
    (requires True)
    (ensures (fun res -> res == List.Tot.partition f l))
= match l with
  | [] -> [], []
  | hd::tl ->
     let l1, l2 = my_partition f tl in
     if f hd
     then hd::l1, l2
     else l1, hd::l2

let rec partition_ext (#a: Type) (f1 f2: (a -> Tot bool)) (l: list a)
: Lemma
  (requires (forall x . f1 x == f2 x))
  (ensures (List.Tot.partition f1 l == List.Tot.partition f2 l))
= match l with
  | [] -> ()
  | hd::tl -> partition_ext f1 f2 tl

let my_bool_of_compare (#a: Type) (f: a -> a -> Tot int) (x: a) (y: a) : Tot bool
= f x y < 0

let rec my_sortWith (#a: Type) (f: (a -> a -> Tot int)) (l:list a)
  : Pure (list a)
    (requires True)
    (ensures (fun res -> res == List.Tot.sortWith f l))
    (decreases (List.Tot.length l))
= match l with
  | [] -> []
  | pivot::tl ->
     let hi, lo = my_partition (my_bool_of_compare f pivot) tl in
     partition_ext (my_bool_of_compare f pivot) (List.Tot.bool_of_compare f pivot) tl;
     List.Tot.partition_length (List.Tot.bool_of_compare f pivot) tl;
     my_append (my_sortWith f lo) (pivot::my_sortWith f hi)

let rec sortWith_ext (#a: Type) (f1 f2: (a -> a -> Tot int)) (l: list a)
: Lemma
  (requires (forall x y . f1 x y == f2 x y))
  (ensures (List.Tot.sortWith f1 l == List.Tot.sortWith f2 l))
  (decreases (List.Tot.length l))
= match l with
  | [] -> ()
  | pivot::tl ->
    partition_ext (List.Tot.bool_of_compare f1 pivot) (List.Tot.bool_of_compare f2 pivot) tl;
    List.Tot.partition_length (List.Tot.bool_of_compare f1 pivot) tl;
    let hi, lo = List.Tot.partition (List.Tot.bool_of_compare f1 pivot) tl in
    sortWith_ext f1 f2 lo;
    sortWith_ext f1 f2 hi

let permute = list atom -> list atom

let my_compare_of_bool (#a:eqtype) (rel: a -> a -> Tot bool) (x: a) (y: a) : Tot int
=   if x `rel` y  then -1
    else if x = y then 0
    else 1
val sort:permute
Steel.Effect.Common.permute