Simple tactic to prove associativity in setoid monoids

SimpleAlgebraTactics.fst

module SimpleAlgebraTactics

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

noeq type setoid_monoid (a: Type) = {
  eq: a -> a -> Type; (* Equivalence relation *)
  op: a -> a -> a;  
  id: a; (* Identity element *)
  (* Equivalence relation properties *)
  refl: x:a -> Lemma(eq x x);
  symm: x:a -> y:a -> eq x y -> Lemma(eq y x);
  trans: x:a -> y:a -> z:a -> Lemma (requires eq x y /\ eq y z) (ensures 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));

  left_congruence: x:a -> y:a -> z:a -> 
                 Lemma (requires eq x y) (ensures eq (op x z) (op y z));
  right_congruence: x:a -> y:a -> z:a -> 
                 Lemma (requires eq x y) (ensures eq (op z x) (op z y));
  (* Identity element *)
  left_identity: x:a -> Lemma (ensures eq (op id x) x);
  right_identity: x:a -> Lemma (ensures eq (op x id) x);  
}

type monoid_mul #a (m: setoid_monoid a) = (mul: (a -> a -> a){mul == m.op})
type monoid_eq #a (m: setoid_monoid a) = (eq: (a -> a -> Type){eq == m.eq})

let left_congruence_lemma #a (m:setoid_monoid a) (eq: monoid_eq m) (m1 m2: monoid_mul m) (x y z: a) (p: unit{eq x y})
  : squash (eq (m1 x z) (m2 y z)) = m.left_congruence x y z
  
let right_congruence_lemma #a (m:setoid_monoid a) (eq: monoid_eq m) (m1 m2: monoid_mul m) (x y z: a) (p: unit{eq x y})
  : squash (eq (m1 z x) (m2 z y)) = m.right_congruence x y z
  
let trans_lemma #a (m: setoid_monoid a) (eq: monoid_eq m) (x y z: a) (p: unit{eq x y /\ eq y z})
  : squash (eq x z) = m.trans x y z
 
let assoc_lemma #a (m:setoid_monoid a) (eq: monoid_eq m) (m1 m2 m3 m4: monoid_mul m) (x y z: a) :
  squash (eq (m1 x (m2 y z)) (m3 (m4 x y) z)) = m.associativity x y z

let term_equal (t1 t2: term) : Tac bool =
  let env = cur_env() in
  let t1_norm = norm_term [primops; iota; zeta; delta_attr ["unfold"]; delta] t1 in
  let t2_norm = norm_term [primops; iota; zeta; delta_attr ["unfold"]; delta] t2 in
  term_eq t1_norm t2_norm

let rec get_app_as_list (t: term) (op1 op2: 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 op1 op2
      else if term_equal f op1 then f::[arg]
      else if term_equal f op2 then f::[arg]
      else (get_app_as_list f op1 op2)@[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) (op1 op2:term) : Tac binary_or_unary_op =
  match inspect (maybe_unsquash_term op) with
  | Tv_App _ _ -> 
                  (match get_app_as_list op op1 op2 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 #t (m: setoid_monoid t) : Tac unit =
    let atype = quote t in
    let mon_mul =  quote m.op in
    let mon_eq = quote m.eq in
    let insp_term term = op_inspect term mon_eq mon_mul in
    match insp_term (maybe_unsquash_term (cur_goal())) with
    | BinOp op lhs rhs ->
        let lhs_app = get_app_as_list lhs mon_eq mon_mul in
        let rhs_app = get_app_as_list rhs mon_eq mon_mul in
        if term_equal op (quote m.eq)
        then begin match (lhs_app, rhs_app) with
          | ([mul1; a; b], [mul2; c; d]) -> if term_equal mul1 mul2 then
          begin match (insp_term a, insp_term b, insp_term c, insp_term 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 m)) in           
              let inspect_assoc = pack (Tv_App inspect_assoc (op, Q_Explicit)) in
              let inspect_assoc = pack (Tv_App inspect_assoc (mul1, Q_Explicit)) in
              let inspect_assoc = pack (Tv_App inspect_assoc (innerLeft, Q_Explicit)) in
              let inspect_assoc = pack (Tv_App inspect_assoc (innerRight, Q_Explicit)) in
              let inspect_assoc = pack (Tv_App inspect_assoc (mul2, 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 
               
              mapply assoc_term 
            | (Atom la, Atom la1, _, Atom ra) ->  fail (term_to_string la1)
            
              //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"
    | Atom wat -> fail ("Top-level op is " ^ term_to_string wat)
    | _ -> fail "Top-level operator is neither binary nor unary"
   
   
let example_proof 
  (a:Type)  
  (m: setoid_monoid a) 
  (mul: a->a->a{mul == m.op})
  (x y z: a)
  : Lemma (m.eq (m.op x (m.op y z)) (m.op (m.op x y) z))
  =
  assert(mul x (m.op y z) `m.eq` (mul (mul x y) z)) by apply_associativity m
 
let print_context () : Tac unit =
   dump "wat"

let apply_transitivity #t (m: setoid_monoid t) (eq: monoid_eq m) (midterm: term) : Tac unit =
  let op_inspect term = op_inspect term (quote m.op) (quote m.eq) 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 not (term_equal op (quote eq)) then fail "Goal is not an eq call"
    else begin
      let trans_lemma_construction = quote (trans_lemma m) in
      let trans_lemma_construction = pack (Tv_App trans_lemma_construction (op, Q_Explicit)) in
      let trans_lemma_construction = pack (Tv_App trans_lemma_construction (lhs, Q_Explicit)) in
      let trans_lemma_construction = pack (Tv_App trans_lemma_construction (midterm, Q_Explicit)) in
      let trans_lemma_construction = pack (Tv_App trans_lemma_construction (rhs, Q_Explicit)) in
      apply trans_lemma_construction
    end
  | _ -> fail "Top-level operator is neither binary nor unary"

let apply_transitivity_nomid #t (eq: t -> t -> Type) (m: setoid_monoid t) : Tac unit =
  let op_inspect term = op_inspect term (quote m.op) (quote m.eq) 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 not (term_is_lookup_equal_to_symbol op eq) then fail "Goal is not an eq call"
    else begin
      let trans_lemma_construction = quote (trans_lemma m) in
      //let trans_lemma_construction = pack (Tv_App trans_lemma_construction (lhs, Q_Explicit)) in
      apply_lemma trans_lemma_construction
    end
  | _ -> fail "Top-level operator is neither binary nor unary"

let transitivity_example_proof #a (m: setoid_monoid a) (x y z: a) 
: Lemma (requires m.eq x y /\ m.eq y z) (ensures m.eq x z) =
  assert(m.eq x z) by apply_transitivity m m.eq (quote y)