Tactic example featuring associativity lemma application

AssociativityTest.fst

module TacSandbox

open FStar.Tactics.V2
open FStar.List
open FStar.List.Tot

noeq type setoid_monoid (a: Type) (eq: a -> a -> Type) = {
  op: a -> a -> a;
  
  (* Equivalence relation properties *)
  refl: x:a -> eq x x;
  symm: x:a -> y:a -> eq x y -> eq y x;
  trans: x:a -> y:a -> z:a -> 
         eq x y -> eq y z -> eq x z;
  
  (* Associativity lemma using the custom equivalence *)
  associativity: x:a -> y:a -> z:a -> 
                 Lemma(eq (op x (op y z)) (op (op x y) z))                  
}

let assoc_lemma (a:Type) (eq: a -> a -> Type) (m:setoid_monoid a eq) (mul: a->a->a{mul == m.op}) (x y z: a) :
  squash (eq (mul x (mul y z)) (mul (mul x y) z)) =
  m.associativity x y z

let rec get_app_as_list (t: term) : Tac (list term) =
  match inspect t with
  | Tv_App f (arg, q) -> 
      if (term_to_string f = "squash") then get_app_as_list arg
      else (get_app_as_list f)@[arg]
  | _ -> [t]

noeq type binary_or_unary_op = 
  | BinOp : (op: term) -> (lhs: term) -> (rhs: term) -> binary_or_unary_op
  | UnOp : (op: term) -> (arg: term) -> binary_or_unary_op
  | Atom : (term) -> binary_or_unary_op

let op_inspect (op: term) : Tac binary_or_unary_op =
  match inspect (maybe_unsquash_term op) with
  | Tv_App _ _ -> 
                  (match get_app_as_list op with
                  | [a1; a2; a3; a4] -> fail ("Too many arguments: (" ^ term_to_string a1 ^ ", " ^ term_to_string a2 ^ ", " ^ term_to_string a3 ^ ", " ^ term_to_string a4 ^ ")")
                  | [o; a1; a2] -> BinOp o a1 a2
                  | [o; a1] -> UnOp o a1
                  | _ -> Atom op)
  | _ -> Atom op

let rec string_concat (sep: string) (lst: list string) : string =
  match lst with
  | [] -> ""
  | [x] -> x
  | x::xs -> x ^ sep ^ (string_concat sep xs)

let term_name_of (t: term) : Tac string =
  match inspect t with
  | Tv_FVar fv -> flatten_name (inspect_fv fv) 
  | Tv_BVar namedv -> bv_to_string namedv
  | Tv_Var v -> term_to_string t
  | _ -> term_to_string t

let symbol_name_of x : Tac string =
  let ty = quote x in
  match inspect ty with
  | Tv_FVar fv -> 
      let name = flatten_name (inspect_fv fv) in
      name
  | Tv_BVar namedv -> bv_to_string namedv
  | Tv_Var v -> term_to_string ty
  | _ -> fail "Could not extract type name" 

let term_is_lookup_equal_to_symbol (t: term) symbol : Tac bool =
  let name_of_symbol = symbol_name_of symbol in
  let name_of_term = term_name_of t in
    name_of_symbol = name_of_term

let rec term_list_to_string (lst: list term) : Tac string =
   match lst with
    | [] -> ""
    | [x] -> term_to_string x
    | x::xs -> term_to_string x ^ ", " ^ (term_list_to_string xs)

let apply_associativity #aty (eq: aty->aty->Type) (m: setoid_monoid aty eq) : Tac unit =
    let atype = quote aty in
    match op_inspect (cur_goal()) with
    | BinOp op lhs rhs ->
        let lhs_app = get_app_as_list lhs in
        let rhs_app = get_app_as_list rhs in
        if term_is_lookup_equal_to_symbol op eq 
        then begin match (lhs_app, rhs_app) with
          | ([mul1; a; b], [mul2; c; d]) -> if term_name_of mul1 = term_name_of mul2 then
          begin match (op_inspect a, op_inspect b, op_inspect c, op_inspect d) with
            | (BinOp innerLeft ia ib, _, _, BinOp innerRight ic id) 
              -> fail "not implemented yet"   
            | (Atom la, BinOp innerLeft ia ib, BinOp innerRight ic id, Atom ra) ->   
              let x = la in
              let y = ia in
              let z = ib in
              let inspect_assoc = (quote (assoc_lemma aty eq m)) in           
              let inspect_assoc = pack (Tv_App inspect_assoc (mul1, Q_Explicit)) in
              let assoc_term = pack (Tv_App inspect_assoc (x, Q_Explicit)) in
              let assoc_term = pack (Tv_App assoc_term (y, Q_Explicit)) in
              let assoc_term = pack (Tv_App assoc_term (z, Q_Explicit)) in
              
              apply assoc_term
 
              //fail (term_to_string la ^ " `" ^ term_to_string mul1 ^ "` (" ^ term_to_string ia ^ " `" ^ term_to_string innerLeft ^ "` " ^ term_to_string ib ^ ")" ^ " vs (" ^ term_to_string ic ^ " `" ^ term_to_string innerRight ^ "` " ^ term_to_string id ^ ") `" ^ term_to_string mul2 ^ "` " ^ term_to_string ra)
            |_ -> fail ("Mul1: " ^ term_to_string mul1 ^ ", Mul2: " ^ term_to_string mul2 ^ ", a: " ^ term_to_string a ^ ", b: " ^ term_to_string b ^ ", c: " ^ term_to_string c ^ ", d: " ^ term_to_string d)
          end else fail "left and right hand operators are not equal"
          | _ -> fail "both hands of the equation should be binary operators"
        end 
        //fail ("Left-hand list is: " ^ term_list_to_string lhs_app ^ ", right-hand list is " ^ term_list_to_string rhs_app)
        else fail "Top-level operator is not the monoid equivalence."
    | UnOp op arg -> fail "Top-level operator is unary"
    | _ -> fail "Top-level operator is neither binary nor unary"
   

(* Example usage *)
let example_proof 
  (a:Type) 
  (eq: a -> a -> Type)
  (m: setoid_monoid a eq) 
  (mul: a->a->a{mul == m.op})
  (x y z: a)
  : Lemma (eq (m.op x (m.op y z)) (m.op (m.op x y) z))
  =
  assert(mul x (mul y z) `eq` (mul (mul x y) z)) by apply_associativity eq m