(* Here we provide an observational equivalence checker for a form of
   first order idealised algol. This is a direct ocaml implementation
   of the ideas presented in _Algorithmic Game Semantics : A Tutorial
   Introduction_ by Samson Abramsky. Further documentation to perhaps
   come. Note in this version of expressions do not have side effects
   (i.e. a;b is only valid if b is of command type) but this could be
   changed by changing the type of seq plus enhancing our operational
   semantic function. *)

(* We firstly need to define tail recursive versions of some List functions.
   We will quite often be using lists to model sets, and as such will often
   use e.g. List.rev_map rather than List.map since the former is tail
   recursive and ordering does not matter since we're representing sets. *)
   
let (@) xs ys = 
  List.rev_append xs ys

let tl_append xs ys =
  List.rev_append (List.rev xs) ys

let rev_concat xss =
  let rec rev_concat acc = function
    | [] -> acc
    | x::xs -> rev_concat (x @ acc) xs in
  rev_concat [] xss

let rev_combine xs ys =
  let rec rev_combine acc xs = function
    | [] -> acc
    | y::ys -> match xs with x::xs -> rev_combine ((x, y)::acc) xs ys in
  rev_combine [] xs ys
  
let product xs ys =
  let rec product ys acc = function
    | x::xs -> product ys (List.rev_map (fun y -> (y,x)) ys @ acc) xs
    | [] -> acc
  in product xs [] ys

let distinct xs = (* we probably now only use this in sorted situations... *)
  let rec distinct ys = function
    | [] -> ys
    | x::xs -> if List.mem x ys then distinct ys xs else distinct (x::ys) xs in
  distinct [] xs

let groupby f xs =
  let rec groupby f acc = function
    | [] -> acc
    | x::xs ->
        if List.exists (fun y -> f x (List.hd y)) acc then
          groupby f (List.rev_map (fun y -> if f x (List.hd y) then x::y else y) acc) xs
        else
          groupby f ([x]::acc) xs
  in groupby f [] xs

(* MODULE TO REPRESENT FINITE AUTOMATA *)

module Ndfa = struct
 
  (* Finite automata over an alphabet *)
  type ('a, 'b) reg_auto =
   { init_state : 'a;
     trans : ('a * 'b * 'a) list;
     fin_states : 'a list }

  (* val from : ('a, 'b) reg_auto -> 'a -> 'a list *)
  let from a s =
      List.rev_map (fun (_,_,y) -> y) (List.filter (fun (x,_,_) -> x = s) a.trans)

  (* val get_states : ('a, 'b) reg_auto -> 'a *)
  (* Returns all accessible states of a given automaton *)
  let get_states a =
      let rec bfs found = function
        | x::xs -> if List.mem x found then bfs found xs else bfs (found @ [x]) (xs @ (from a x))
        | [] -> found in
      bfs [] [a.init_state]

  let get_alpha a =
      distinct (List.rev_map (fun (_,y,_) -> y) a.trans)
            
  (* val clean_auto : ('a, 'b) reg_auto -> ('a, 'b) reg_auto * int *)
  (* Removes an automaton of all inaccessible states, and returns the number of such states *)
  let clean_auto auto =
      let states = get_states auto in
      ({  init_state = auto.init_state;
          trans = List.filter (fun (s,_,_) -> List.mem s states) auto.trans;
          fin_states = List.filter (fun s -> List.mem s states) auto.fin_states},
      List.length states)

  (* val sort_auto : ('a, 'b) reg_auto -> (int, 'b) reg_auto *)
  (* Relabels the states of an automaton to be integers in a sensible order *)
  let sort_auto auto =
      let states = get_states auto in
      let rec numb_states n acc = function (* and renumber from 0 *)
        | s::ss -> numb_states (n+1) ((s,n)::acc) ss
        | [] -> acc in
      let n_states = numb_states 0 [] states in
      let change_trans (s,a,r) = (List.assoc s n_states, a, List.assoc r n_states) in
      { init_state = List.assoc auto.init_state n_states;
        trans = List.rev_map change_trans auto.trans;
        fin_states = List.rev_map (fun s -> List.assoc s n_states) auto.fin_states }

  (* delta_star : ('a, 'b) reg_auto -> 'a -> 'b list -> 'a list *)
  (* Follows a path in an automaton from an initial state *)
  let rec delta_star a q = function
    | [] -> [q]
    | l::ls ->
          let start_trans = List.filter (fun (x,y,_) -> x = q && y = l) a.trans in
          let new_states = List.rev_map (fun (_,_,z) -> z) start_trans in
          rev_concat (List.rev_map (fun x -> delta_star a x ls) new_states)

  (* val accept_word : ('a, 'b) reg_auto -> 'b list -> bool *)
  (* Determines whether a given automaton accepts a word *)
  let accept_word a lets =
      let fin_states = delta_star a a.init_state lets in
      let passed = List.rev_map (fun x -> List.mem x a.fin_states) fin_states in
      List.fold_left (||) false passed

  (* val accepts_nothing : ('a, 'b) reg_auto -> bool *)
  (* Determines whether there is a word that this automaton accepts *)
  let accepts_nothing a =
      let rec bfs found = function
        | x::xs ->
              if List.mem x a.fin_states then false
              else if List.mem x found then bfs found xs
              else bfs (x::found) (xs @ (from a x))
        | [] -> true in
      bfs [] [a.init_state]
    
  (* val example_word : ('a, 'b) reg_auto -> 'b list option *)
  (* Returns an example of a word an automaton accepts, if it accepts any *) 
  let example_word a =
      let from s psf =
          let some_trans = List.filter (fun (x,_,_) -> x = s) a.trans in
          List.rev_map (fun (_,y,z) -> (z, tl_append psf [y])) some_trans in
      let rec bfs found = function
        | (x,p)::xs ->
              if List.mem x a.fin_states then Some p
              else if List.mem x found then bfs found xs
              else bfs (x::found) (xs @ (from x p))
        | [] -> None in
      bfs [] [(a.init_state, [])]

  (* delta_from and determinise can deal with automaton with epsilon-transitions.
  These eps-transitions are input as a seperate parameter. If a has an eps-link to
  b then effectively all arrows from b need to be "copied" to also come from a. *)
  
  (* delta_from : (a', 'b) reg_auto -> ('a * 'a) list -> 'a -> ('b * 'a) list *)
  let delta_from a etrans q =
      let rec eps_from a q xs =
          if List.mem q xs then [q] else
          let etrans = List.filter (fun (x,_) -> x = q) etrans in
          let estates = List.rev_map (fun (x,y) -> y) etrans in
          let efromrest = List.rev_map (fun y -> eps_from a y (q::xs)) estates in
          q::(rev_concat efromrest) in
      let from_a_noeps a q =
          List.rev_map (fun (_,y,z) -> (y,z)) (List.filter (fun (x,_,_) -> x = q) a.trans) in
      distinct (rev_concat (List.rev_map (fun q -> from_a_noeps a q) (eps_from a q [])))
  
  (* determinise : ('a, 'b) reg_auto -> ('a * 'a) list -> (int, 'b) reg_auto * int *)    
  let determinise a etrans =
      let contains_fs ps = List.exists (fun f -> List.mem f ps) a.fin_states in
      let rec bfs trans fstates = function
        | x::xs ->
            if List.mem x (List.rev_map (fun (x,y,z) -> x) trans) then bfs trans fstates xs
            else
            let ys = List.rev_map (fun q -> delta_from a etrans q) x in
            let ys = rev_concat ys in
            let ys = groupby (fun (y1,z1) (y2,z2) -> y1 = y2) ys in
            let ys = List.rev_map (function (y,z)::zs -> (y, z::(List.rev_map (fun (y,z) -> z) zs))) ys in
            let ys = List.rev_map (fun (y,zs) -> (x,y,List.fast_sort (fun x y -> if x < y then -1 else 1) (distinct zs))) ys in
            bfs (ys @ trans) (if contains_fs x then x::fstates else fstates) (xs @ (distinct (List.rev_map (fun (x,y,z) -> z) ys)))
        | [] -> (trans,fstates) in
      let (trans, fstates) = bfs [] [] [[a.init_state]] in
      let auto = { init_state = [a.init_state]; trans = trans; fin_states = distinct fstates } in
      (sort_auto auto, List.length (get_states auto))

  (* increase_states : (int, 'b) reg_auto -> int -> (int, 'b) reg_auto *)
  let increase_states auto n =
    { init_state = auto.init_state + n;
      trans = List.rev_map (fun (x,y,z) -> (x+n, y, z+n)) auto.trans;
      fin_states = List.rev_map ((+) n) auto.fin_states }
                      
end

(* MODULE TO REPRESENT REGULAR EXPRESSIONS *)

module RegExp = struct

  (* Regular expressions over an alphabet *)
  (* These are extended but can all be reduced to NDFAs *)
  type 'b reg_exp =
    | Empty
    | Eps
    | Char of 'b
    | Alt of 'b reg_exp * 'b reg_exp
    | Star of 'b reg_exp
    | Seq of 'b reg_exp * 'b reg_exp
    | Both of 'b reg_exp * 'b reg_exp
    | Hom of ('b -> 'b list) * 'b reg_exp
    | InvHom of ('b -> 'b list) * 'b list * 'b reg_exp
    | Sigma of int * (int -> 'b reg_exp)
    | Comp of 'b list * 'b reg_exp 

  (* val seq : 'b reg_exp list -> 'b reg_exp *)
  let rec seq = function
    | [] -> Eps
    | [x] -> x
    | (x::xs) -> Seq(x, seq xs)

  (* val lit : 'b list -> 'b reg_exp *)
  let rec lit = function
    | [] -> Eps
    | [x] -> Char x
    | x::xs -> Seq(Char x, lit xs)

  (* val alt : 'b reg_exp list -> 'b reg_exp *)    
  let rec alt = function
    | [] -> Empty
    | [x] -> x
    | x::xs -> Alt(x, alt xs)
    
  (* val rem_sums : 'b reg_exp -> 'b reg_exp *)
  (* Converts the Sigma operator into nested Alts *)
  let rec rem_sums = function
    | Eps -> Eps | Empty -> Empty | Char n -> Char n 
    | Alt (r1,r2) -> Alt (rem_sums r1, rem_sums r2)
    | Star r -> Star (rem_sums r) | Both (r1,r2) -> Both (rem_sums r1, rem_sums r2)
    | Seq (r1, r2) -> Seq (rem_sums r1, rem_sums r2)
    | Hom (f,r) -> Hom (f, rem_sums r) | InvHom (f,l,r) -> InvHom (f, l, rem_sums r)
    | Comp (a, r) -> Comp (a, rem_sums r)
    | Sigma (n, f) ->
          let rec to_regular (n, f) = match n with
            | 0 -> Empty
            | 1 -> rem_sums (f n)
            | n -> Alt (rem_sums (f n), to_regular (n-1, f)) in
          to_regular (n, f)
    
  open Ndfa (* Regular expressions are equivalent to Ndfas *)

  (* val compile_regexp : 'b reg_exp -> (int, 'b) reg_auto *)
  (* Compiles a regular expression to an automaton *)
  let compile_regexp r =
      let rec compile_regexpr = function
        | Empty -> ({ init_state = 0; trans = []; fin_states = [] }, 1)
        | Eps -> ({ init_state = 0; trans = []; fin_states = [0] }, 1)
        | Char b -> ({ init_state = 0; trans = [(0,b,1)]; fin_states = [1]}, 2)
        | Alt (r1,r2) ->
              let (a1,n1) = compile_regexpr r1 in
              let (a2,n2) = compile_regexpr r2 in
              let a2 = increase_states a2 n1 in
              let alter_trans (x,y,z) =
                  if (x = a2.init_state) then (a1.init_state,y,z - 1) else (x - 1,y,z - 1) in
              let add_fin = if List.mem a2.init_state a2.fin_states then [a1.init_state] else [] in
              let fin_states = add_fin @ a1.fin_states @ (List.rev_map pred a2.fin_states) in
              ({ init_state = 0;
                 trans = a1.trans @ (List.rev_map alter_trans a2.trans);
                 fin_states = fin_states}, n1 + n2 - 1)
        | Star r ->
              let (a,n) = compile_regexpr r in
              let start_trans = List.filter (fun (x,y,z) -> x = a.init_state) a.trans in
              let link_final fs = List.rev_map (fun (x,y,z) -> (fs,y,z)) start_trans in
              let extra_trans = rev_concat (List.rev_map link_final a.fin_states) in
              ({ init_state = a.init_state;
                 trans = a.trans @ extra_trans;
                 fin_states = a.init_state::a.fin_states}, n)
        | Seq (r1,r2) ->
              if r1 = Eps then compile_regexpr r2 else
              if r2 = Eps then compile_regexpr r1 else
              let (a1,n1) = compile_regexpr r1 in
              let (a2,n2) = compile_regexpr r2 in
              let a2 = increase_states a2 n1 in
              let link_finals l s = List.rev_map (fun fs -> (fs,l,s)) a1.fin_states in
              let alter_trans (x,y,z) =
                  if (x = a2.init_state) then link_finals y (z - 1) else [(x - 1,y,z - 1)] in
              let comb_trans = a1.trans @ (rev_concat (List.rev_map alter_trans a2.trans)) in
              let add_fin = if List.mem a2.init_state a2.fin_states then a1.fin_states else [] in
              let fin_states = add_fin @ (List.rev_map pred a2.fin_states) in
              ({ init_state = a1.init_state;
                 trans = comb_trans;
                 fin_states = fin_states }, n1 + n2 - 1)
        | InvHom (f,lang,r) ->
              let (a, n) = compile_regexpr r in
              let states = get_states a in
              let trans = List.rev_map (fun (q,l) -> List.rev_map (fun x -> (q, l, x)) (delta_star a q (f l))) (product states lang) in
              let trans = rev_concat trans in
              let new_auto = { init_state = a.init_state;
                 trans = trans;
                 fin_states = a.fin_states } in
              let (cleaned_auto, nrem) = clean_auto new_auto in
              (sort_auto cleaned_auto, nrem)
        | Hom (f, r) ->
              let (a, n) = compile_regexpr r in
              let alter_trans (s1,al,s2) =
                  let letters = f al in
                  let rec interim n acc = function
                    | [] -> (acc, [(s1,0),(s2,0)])
                    | [x] -> (((s1,n),x,(s2,0))::acc,[])
                    | x::xs -> interim (n+1) (((s1,n),x,(s1,n+1))::acc) xs in
                  interim 0 [] letters in
              let transpairs = List.rev_map alter_trans a.trans in
              let remdups = List.fold_left (fun y x -> if List.mem x y then y else y @ [x]) [] in
              let etrans = remdups (rev_concat (List.rev_map snd transpairs)) in
              let new_auto = { init_state = (a.init_state, 0);
                 trans = rev_concat (List.rev_map fst transpairs);
                 fin_states = List.rev_map (fun a -> (a, 0)) a.fin_states } in
              determinise new_auto etrans
              (* let (cleaned_auto, nrem) = clean_auto new_auto in
              (sort_auto cleaned_auto, nrem) *)
        | Both (r1, r2) ->
              let (a1, n1) = compile_regexpr r1 in
              let (a2, n2) = compile_regexpr r2 in
              let joined_trans =
                  List.filter (fun ((s1,a,s2),(r1,b,r2)) -> a = b) (product a1.trans a2.trans) in
              let output_trans =
                  List.rev_map (fun ((s1,a,s2),(r1,b,r2)) -> ((s1,r1), a, (s2,r2))) joined_trans in
              let (prod_auto, nprod) = ({
                  init_state = (a1.init_state,a2.init_state);
                  trans = output_trans;
                  fin_states = product a1.fin_states a2.fin_states }, n1 * n2) in
              let (cleaned_auto, nrem) = clean_auto prod_auto in
              (sort_auto cleaned_auto, nrem)
        | Comp (lang,r) ->
              let (a, n) = determinise (fst (compile_regexpr r)) [] in
              let states = get_states a in
              let no_trans = List.filter (fun (s,l) -> delta_star a s [l] = []) (product states lang) in
              let extra_trans = List.rev_map (fun (s,l) -> (s,l,n)) no_trans in
              let more_trans = List.rev_map (fun l -> (n,l,n)) lang in
              let fin_states = n :: (List.filter (fun s -> not (List.mem s a.fin_states)) states) in
              ({ init_state = a.init_state;
                 trans = a.trans @ extra_trans @ more_trans;
                 fin_states = fin_states }, n)
         | Sigma (n, f) as r ->
               compile_regexpr (rem_sums r) in
      fst (compile_regexpr r)
      
  (* val match_exp : 'b regexp -> 'b list -> bool *)
  (* Determines whether a word satisfies a regular expression *)
  let match_exp w r = accept_word (compile_regexp r) w

  (* val is_empty : 'b regexp -> bool *)
  (* Determines whether a regular expression is empty *)
  let is_empty r = accepts_nothing (compile_regexp r)

  (* val get_lang : 'b reg_exp -> 'b list *)
  let get_lang bs =
    let rec get_lang = function
      | Empty | Eps -> []
      | Char s -> [s]
      | Seq(s1,s2) | Alt(s1,s2) | Both(s1,s2) -> get_lang s1 @ (get_lang s2)
      | Star s -> get_lang s | Comp(l,s) -> get_lang s @ l
      | Sigma _ as t -> get_lang (rem_sums t)
      | Hom (f, r) -> rev_concat (List.rev_map f (get_lang r))
      | InvHom (f,l,r) -> l
    in distinct (get_lang bs)
    
  (* val contains : 'b list -> 'b regexp -> 'b regexp -> bool *)
  (* Determines whether, over a given alphabet, one regexp is contained within another *)
  let contains r1 r2 = is_empty (Both(r1, Comp(get_lang r1,r2)))

  (* val equal : 'b list -> 'b regexp -> 'b regexp -> bool *)
  (* Determines whether, over a given alphabet, two regular expressions represent the same set *)
  let equal r1 r2 = contains r1 r2 && (contains r2 r1)

  (* val example_word : 'b regexp -> Some ('b list) *)
  (* Returns a word satisfying the given regular expression, if such a word exists. *)
  let example_word r = example_word (compile_regexp r)

end

(* MODULE TO REPRESENT FIRST ORDER IDEALISED ALGOL *)
(* CONTAINS SYNTAX, TYPING RULES AND OPERATIONAL SEMANTICS *)

module Language = struct

  type btype = Com | Nat | Bool | VarN
  type ttype = Basic of btype | Arrow of btype * ttype

  (* val flatten : ttype -> btype list *)
  let rec flatten = function
    | Basic b -> [b]
    | Arrow (b, rest) -> b :: (flatten rest)

  type ident = string
  type context = (ident * ttype) list

  type constant = Skip | While | KSeq | Cond | KNat of int | KBool of bool | Add | Subtr | Not | Deref | Store | And | Equals | Less | Even
  
  let const_type = function
    | Skip -> Basic Com
    | While -> Arrow(Bool, Arrow(Com, Basic Com))
    | KSeq -> Arrow(Com, Arrow(Com, Basic Com))
    | Cond -> Arrow(Bool, Arrow(Com, Arrow(Com, Basic Com)))
    | KNat _ -> Basic (Nat)
    | KBool _ -> Basic (Bool)
    | Add | Subtr -> Arrow (Nat, Arrow(Nat, Basic Nat))
    | And -> Arrow (Bool, Arrow (Bool, Basic Bool))
    | Equals -> Arrow (Nat, Arrow(Nat, Basic Bool))
    | Less -> Arrow (Nat, Arrow(Nat, Basic Bool))
    | Not -> Arrow (Bool, Basic Bool)
    | Deref -> Arrow (VarN, Basic Nat)
    | Store -> Arrow (VarN, Arrow (Nat, Basic Com))
    | Even -> Arrow (Nat, Basic Bool)

  type term = Var of ident | Abs of ident * btype * term | App of term * term | Const of constant | New of ident * term

  (* val get_type : context -> term -> ttype *)
  let rec get_type ctx = function
    | Const c -> const_type c
    | Var x -> List.assoc x ctx
    | Abs (x, t, b) -> Arrow (t, get_type ((x,Basic t)::ctx) b)
    | New (x, b) -> get_type ((x,Basic VarN)::ctx) b
    | App (t1, t2) ->
          match (get_type ctx t1, get_type ctx t2) with
          | (Arrow(b1,t), Basic b2) when b1 = b2 -> t
          | _ -> failwith "The term input is untypable."

  (* val subs : term -> ident -> term -> term *)
  let rec subs m x n = match m with
    | Var y when x = y -> n
    | Var y -> Var y
    | Const k -> Const k
    | Abs (y,bt,tm) when y = x -> Abs (y, bt, tm)
    | Abs (y, bt, tm) -> Abs (y, bt, subs tm x n)
    | App (t1, t2) -> App (subs t1 x n, subs t2 x n)
    | New (y,tm) when y = x -> New (y, tm)
    | New (y,tm) -> New (y, subs tm x n)

  type mem = (ident * term) list

  type value = VNat of int | VBool of bool | VUnit

  (* val reduce : term -> value *)
  (* since seq : com * com -> com, not com * nat -> nat we know expressions
     (i.e. terms of type nat) cannot change the store *)
  let reduce t = 
    let rec reduce mem = function
      | Var x -> failwith "Cannot reduce open term."
      | Abs (x,bt,tm) -> failwith "Term not of ground type."
      | Const (KNat n) -> (mem, VNat n)
      | Const (KBool b) -> (mem, VBool b)
      | Const (Skip) -> (mem, VUnit)
      | App (App (Const KSeq, t1), t2) -> let (mem, _) = reduce mem t1 in reduce mem t2
      | App (App (App (Const Cond, t1), t2), t3) -> (match snd (reduce mem t1) with VBool true -> reduce mem t2 | _ -> reduce mem t3)
      | App (App (Const While, t1), t2) as t -> (match snd (reduce mem t1) with VBool true -> let (mem, _) = reduce mem t2 in reduce mem t | _ -> (mem, VUnit))
      | App (Const Even, t1) -> (match snd (reduce mem t1) with VNat n -> (mem, VBool (n mod 2 = 0)) | _ -> failwith "no")
      | App (App (Const Add, t1), t2) ->
            let ((_, t1),(_,t2)) = (reduce mem t1, reduce mem t2) in
            (match (t1,t2) with (VNat n, VNat m) -> (mem,VNat (n + m)) | _ -> failwith "Invalid")
      | App (App (Const Subtr, t1), t2) ->
            let ((_, t1),(_,t2)) = (reduce mem t1, reduce mem t2) in
            (match (t1,t2) with (VNat n, VNat m) -> (mem,VNat (n - m)) | _ -> failwith "Invalid")
      | App (Const Not, arg) ->
            let (_, t1) = reduce mem arg in
            (match t1 with VBool n -> (mem,VBool (not n)) | _ -> failwith "Invalid")
      | App (App (Const And, t1), t2) ->
            let ((_, t1),(_,t2)) = (reduce mem t1, reduce mem t2) in
            (match (t1,t2) with (VBool n, VBool m) -> (mem,VBool (n && m)) | _ -> failwith "Invalid")
      | App (App (Const Less, t1), t2) ->
            let ((_, t1),(_,t2)) = (reduce mem t1, reduce mem t2) in
            (match (t1,t2) with (VNat n, VNat m) -> (mem,VBool (n < m)) | _ -> failwith "Invalid")
      | App (App (Const Equals, t1), t2) ->
            let ((_, t1),(_,t2)) = (reduce mem t1, reduce mem t2) in
            (match (t1,t2) with (VNat n, VNat m) -> (mem,VBool (n < m)) | _ -> failwith "Invalid")
      | App (Const Deref, Var x) -> (mem, List.assoc x mem)
      | App (App (Const Store, Var x), bd)  -> let (_,v) = reduce mem bd in (((x,v)::mem),VUnit)
      | App (Abs (x, bt, bd), arg) -> reduce mem (subs bd x arg)
      | New (x, bd) -> reduce ((x, VNat 0)::mem) bd
      | _ -> failwith "Invalid type" in
    ignore (get_type [] t); snd (reduce [] t)

end

(* MODULE TO PROVIDE DENOTATION OF TERMS FROM THE ABOVE LANGUAGE *)
(* We represent a term as a strategy, i.e. a regular expression over
   moves of the appropriate game. *)
   
module Strategiser = struct
  open Language
  open RegExp

  (* The alphabet to which all of our moves will belong *)
  type alpha =
    | Q | An of int | Ab of bool | Run | Done | Read | Write of int | Value of int | Ok

  (* val mapn : (int -> 'a) -> int -> 'a list *)
  let mapn f n =
    let rec mapn f acc = function
      | 0 -> acc
      | n -> mapn f ((f n)::acc) (n-1) in
    mapn f [] n

  (* Note we can only represent finite games/alphabets. Since we have a natural number type that is in principle
     unbounded, we must put a finite limit on it: see Samson's paper for details, but for decidability (reducibility
     to regular expressions) we can only deal with a finitary portion of the language *)

  (* val questions : btype -> int -> alpha list *)
  (* Returns the moves that are questions of a given type *)
  let rec questions n = function
    | Com -> [Run]
    | Nat | Bool -> [Q]
    | VarN -> [Read] @ (mapn (fun n -> Write (n - 1)) n)

  (* val answers : btype -> int -> alpha list *)
  (* Returns the moves that are answers of a given type *)
  let rec answers n = function
    | Com -> [Done]
    | Nat -> mapn (fun n -> An (n - 1)) n
    | Bool -> [Ab true; Ab false]
    | VarN -> [Ok] @ (mapn (fun n -> Value (n - 1)) n)

  (* We now present the types for moves of the whole game, i.e. moves in the context. These consists
     of moves in the subgames above, together with a position in the tableaux (i.e. the tagging in the
     disjoint union. *)
  type move = Ctx of ident * int * alpha | Trm of int * alpha
            | LCtx of ident * int * alpha | LTrm of int * alpha (* used only for locus *)
  
  (* Moves on the left of the turn-style are represented by Ctx(x,n,a)
     where x is the free variable in the term refering to that subgame,
     n is the position in that subgame (since this variable can represent
     a first-order function) and a the move itself.
     Moves on the right of the turn-style are represented by Trm(n,a)
     giving the type of move and where it is played.


   e.g. the strategies
      [[f]] =                    [[Var x]] =
   ... |- N -o N                ... N -o N |- N -o N
              q                          q
         q                          n
         n                          n
              n                          q

  are represented by [Trm(2,Q);Trm(1;Q);Trm(1,An n);Trm(2,An n)] and
  [Ctx(x,2,Q);Ctx(x,1,Q);Ctx(x,1,An n);Ctx(x,2,An n)] respectively.

  Remember the "context" here is considered part of the term (since we're in
  fact dealing with typed terms = typing judgements, but we can work out the
  actual term type ourselves. LCtx and LTrm are only used internally (for
  interaction sequences) and will never appear in an output automaton. *)

  (* val moves : int -> ttype -> move list *)
  (* All of the moves of a given type *)
  let moves n t =
      let types = flatten t in
      let alphas = List.rev_map (fun t -> questions n t @ (answers n t)) types in
      let rec ints xs = function
        | 0 -> xs
        | n -> ints (n::xs) (n-1) in
      let mvs = rev_combine (ints [] (List.length types)) (List.rev alphas) in
      let alphs = List.rev_map (fun (m,al) -> List.rev_map (fun a -> Trm (m,a)) al) mvs in
      rev_concat alphs

  (* val all_moves : int -> context -> ttype -> move list *)
  let all_moves n ctx t =
      let convert_ctx_member (x,typ) =
          List.rev_map (function Trm (m,a) -> Ctx (x,m,a) | _ -> failwith "no") (moves n typ) in
      rev_concat (List.rev_map convert_ctx_member ctx) @ moves n t

  (* val den_const : int -> constant -> move reg_exp *)
  let den_const n = function
    | KNat n -> lit [Trm(1,Q);Trm(1,An n)]
    | KBool n -> lit [Trm(1,Q);Trm(1,Ab n)]
    | Skip -> lit [Trm(1,Run);Trm(1,Done)]
    | KSeq -> lit [Trm(3,Run);Trm(1,Run);Trm(1,Done);Trm(2,Run);Trm(2,Done);Trm(3,Done)]
    | Cond ->
          seq [lit [Trm(4,Run);Trm(1,Q)];
               Alt (lit [Trm(1,Ab true); Trm(2,Run); Trm(2,Done); Trm(4,Done)],
                    lit [Trm(1,Ab false); Trm(3,Run); Trm(3,Done); Trm(4,Done)])]
    | While ->
          seq [lit [Trm(3,Run)];
               Star (lit [Trm(1,Q);Trm(1,Ab true);Trm(2,Run);Trm(2,Done)]);
               lit [Trm(1,Q);Trm(1,Ab false);Trm(3,Done)]]
    | Add ->
         seq [lit [Trm(3,Q);Trm(1,Q)];
              Sigma(n, (fun k -> seq[lit[Trm(1,An (k - 1));Trm(2,Q)];
                  Sigma(n, (fun m -> lit[Trm(2,An (m - 1));Trm(3, An(k+m-2))]))]
              ))]
    | Subtr ->
         seq [lit [Trm(3,Q);Trm(1,Q)];
              Sigma(n, (fun k -> seq[lit[Trm(1,An (k-1));Trm(2,Q)];
                  Sigma(n, (fun m -> lit[Trm(2,An (m-1));Trm(3, An(k-m))]))]
              ))]
    | Equals ->
         seq [lit [Trm(3,Q);Trm(1,Q)];
              Sigma(n, (fun k -> seq[lit[Trm(1,An (k - 1));Trm(2,Q)];
                  Sigma(n, (fun m -> lit[Trm(2,An (m - 1));Trm(3, Ab(k = m))]))]
              ))]
    | Less ->
         seq [lit [Trm(3,Q);Trm(1,Q)];
              Sigma(n, (fun k -> seq[lit[Trm(1,An (k-1));Trm(2,Q)];
                  Sigma(n, (fun m -> lit[Trm(2,An (m-1));Trm(3, Ab(k < m))]))]
              ))]
    | Not ->
         seq [lit [Trm(2,Q);Trm(1,Q)];
              Alt(lit[Trm(1,Ab true);Trm(2, Ab false)], 
                  lit[Trm(1,Ab false);Trm(2, Ab true)])]
    | And ->
         seq [lit [Trm(3,Q);Trm(1,Q)];
              alt [lit[Trm(1,Ab false);Trm(2,Q);Trm(2, Ab false);Trm(3,Ab false)];
                   lit[Trm(1,Ab false);Trm(2,Q);Trm(2, Ab true);Trm(3,Ab false)];
                   lit[Trm(1,Ab true);Trm(2,Q);Trm(2, Ab false);Trm(3,Ab false)];
                   lit[Trm(1,Ab true);Trm(2,Q);Trm(2, Ab true);Trm(3,Ab true)]]]
    | Store -> seq [lit [Trm(3,Run);Trm(2,Q)];
              Sigma (n, (fun n -> lit [Trm(2, An (n - 1)); Trm(1, Write (n - 1))]));
              lit [Trm(1, Ok); Trm(3, Done)]]
    | Deref -> seq [lit [Trm(2,Q);Trm(1,Read)];
              Sigma (n, (fun n -> lit [Trm(1, Value (n - 1)); Trm(2, An (n - 1))]))]
    | Even -> seq [lit [Trm(2,Q);Trm(1,Q)]; Sigma (n, (fun n ->
              lit [Trm(1, An (n - 1)); Trm(2,Ab ((n-1) mod 2 = 0))]))]
                   
  (* val den_sem : int -> context -> term -> move reg_exp *)
  (* Denotational semantics of a typed term as a strategy, represented by a regular
     expression over the alphabet of moves. *)
  let den_sem n ctx t =
   let rec den_sem n ctx = function
    | Const k -> den_const n k
    | Abs (x, bt, bd) ->
          let table = function
            | Ctx (y,_,a) when y = x -> [Trm (1,a)]
            | Trm (n,a) -> [Trm (n+1, a)]
            | t -> [t] (*should not happen*) in
          Hom(table, den_sem n ((x,Basic bt)::ctx) bd)
    | Var x as t -> (* Copycat. See Samson's paper for details. *)
          let pl n a = Ctx(x,n,a) in
          let pr n a = Trm(n,a) in
          let types = flatten (get_type ctx t) in
          let final_dest = List.nth types (List.length types - 1) in
          let final_qns = questions n final_dest in
          let ql = List.length final_qns in
          let final_ans = answers n final_dest in
          let al = List.length final_ans in
          let r i =
              let typ = List.nth types (i-1) in
              let qns = questions n typ in
              let ans = answers n typ in
              let ql = List.length qns in
              let al = List.length ans in
              seq
              [ Sigma(ql, (fun n -> lit
                  [pl i (List.nth qns (n-1));
                   pr i (List.nth qns (n-1))]));
                Sigma(al, (fun n -> lit
                  [pr i (List.nth ans (n-1));
                   pl i (List.nth ans (n-1))]))] in
          let k = List.length types - 1 in
          let ris = Sigma(k, (fun i -> r i)) in
          seq
          [ Sigma(ql, (fun n -> lit
              [pr (List.length types) (List.nth final_qns (n-1));
               pl (List.length types) (List.nth final_qns (n-1))]));
            Star(ris);
            Sigma(al, (fun n -> lit
              [pl (List.length types) (List.nth final_ans (n-1));
               pr (List.length types) (List.nth final_ans (n-1))]))]
    | App (t1, t2) ->
          let (type1, type2) = (get_type ctx t1, get_type ctx t2) in
          let ctx2lang = all_moves n ctx type2 in
          let lang1 = moves n type1 in
          let dupcon = List.rev_map (function Ctx(x,n,a) -> [Ctx(x,n,a);LCtx(x,n,a)] | t -> [t]) ctx2lang in
          let dupcon = rev_concat dupcon in
          let movetm1 = List.rev_map (function Trm(1,a) -> [] | Trm(n,a) -> [LTrm(n - 1, a)] | _ -> failwith "no") lang1 in
          let mid_lang = dupcon @ (rev_concat movetm1) in
          let out2 = function
            | Ctx(x,n,a) -> [Ctx(x,n,a)]
            | LCtx(x,n,a) -> []
            | Trm(1,a) -> [Trm(1,a)]
            | Trm(n,a) -> []
            | LTrm(n,a) -> [] in
          let out1 = function
            | Ctx(x,n,a) -> []
            | LCtx(x,n,a) -> [Ctx(x,n,a)]
            | Trm(n,a) -> [Trm(n,a)]
            | LTrm(n,a) -> [Trm(n+1,a)] in
          let out3 = function
            | Ctx(x,n,a) | LCtx(x,n,a) -> [Ctx(x,n,a)]
            | Trm(n,a) -> [] | LTrm(n,a) -> [Trm(n,a)] in
          let d1 = den_sem n ctx t1 in
          let d2 = den_sem n ctx t2 in
          Hom(out3, Both(InvHom(out2, mid_lang, Star d2),InvHom(out1, mid_lang, d1)))
    | New (x, bd) ->
          let writeseq n = seq [lit [Trm(1,Write n); Trm(1,Ok)];
                               Star(lit [Trm(1,Read); Trm(1,Value n)])] in
          let cell = seq [Star (lit [Trm (1, Read); Trm (1, Value 0)]);
                          Star(Sigma (n, (fun n -> writeseq (n - 1))))] in
          let newctx = ((x,Basic VarN)::ctx) in
          let lang = all_moves n newctx (get_type newctx bd) in
          let out1 = function
           | Ctx(y,n,a) when y = x -> [Trm(n,a)]
           | t -> [] in
          let out2 = function
           | Ctx(y,n,a) when y = x -> []
           | t -> [t] in
          let bd_sem = den_sem n newctx bd in
          Hom(out2, Both(bd_sem, InvHom(out1,lang,cell))) in
   ignore (get_type ctx t); den_sem n ctx t

end

(* MODULE FOR PARSING THE SYNTAX OF THE LANGUAGE *)
(* Fundamentally a mapping from strings to terms *)

module LangParse = struct

  open Genlex
  open Language
  
  let keywds = ["while"; "do"; "done"; "skip"; "if"; "then"; "else";
                "true"; "false"; "+"; "-"; "~"; "&"; "="; "<"; "even";
                "!"; "set"; "="; "new"; "in";
                "fun"; "@"; 
                ":"; "->"; "com"; "nat"; "bool"; "var"; ";";
                "("; ")"; "-"; ".";
                "omega"]
                
  let to_ident x = x
  let lexer = make_lexer keywds
  
  let parse_bt = parser
    | [< 'Kwd "com" >] -> Com
    | [< 'Kwd "nat" >] -> Nat
    | [< 'Kwd "bool" >] -> Bool
    | [< 'Kwd "var" >] -> VarN
    
  let rec parse_type = parser
    | [< bt = parse_bt; typ = parse_rest bt >] -> typ
  and parse_rest bt = parser
    | [< 'Kwd ";" >] -> Basic bt
    | [< 'Kwd "->"; cod = parse_type >] -> Arrow(bt, cod)
    
  let rec parse_expr = parser
    | [< 'Kwd "if"; cond = parse_expr; 'Kwd "then"; thn = parse_expr; 'Kwd "else"; els = parse_expr >] -> App(App(App(Const Cond, cond),thn), els)
    | [< 'Kwd "skip" >] -> Const Skip
    | [< 'Kwd "while"; cond = parse_expr; 'Kwd "do"; bd = parse_expr; 'Kwd "done" >] -> App(App(Const While, cond), bd)
    | [< 'Kwd "fun"; 'Ident x; 'Kwd ":"; typ = parse_bt; 'Kwd "->"; expr = parse_expr >] -> Abs(to_ident x, typ, expr)
    | [< 'Kwd "new"; 'Ident x; 'Kwd "in"; expr = parse_expr >] -> New(to_ident x, expr)    
    | [< 'Kwd "("; expr = parse_expr; 'Kwd ")" >] -> expr
    | [< 'Int n >] -> Const (KNat n)
    | [< 'Kwd "true" >] -> Const (KBool true)
    | [< 'Kwd "false" >] -> Const (KBool false)
    | [< 'Kwd "+"; expr1 = parse_expr; expr2 = parse_expr >] -> App(App(Const Add, expr1), expr2)
    | [< 'Kwd "-"; expr1 = parse_expr; expr2 = parse_expr >] -> App(App(Const Subtr, expr1), expr2)
    | [< 'Kwd "~"; expr = parse_expr >] -> App(Const Not, expr)
    | [< 'Kwd "even"; expr = parse_expr >] -> App(Const Even, expr)
    | [< 'Kwd "&"; expr1 = parse_expr; expr2 = parse_expr >] -> App(App(Const And, expr1), expr2)
    | [< 'Kwd "="; expr1 = parse_expr; expr2 = parse_expr >] -> App(App(Const Equals, expr1), expr2)
    | [< 'Kwd "<" ; expr1 = parse_expr; expr2 = parse_expr >] -> App(App(Const Less, expr1), expr2)
    | [< 'Kwd "!"; expr = parse_expr >] -> App(Const Deref, expr)
    | [< 'Ident x >] -> Var (to_ident x)
    | [< 'Kwd "do"; expr1 = parse_expr; 'Kwd "then"; expr2 = parse_expr >] -> App(App(Const KSeq, expr1), expr2)
    | [< 'Kwd "@"; expr1 = parse_expr; expr2 = parse_expr >] -> App(expr1, expr2)
    | [< 'Kwd "set"; expr1 = parse_expr; 'Kwd "="; expr2 = parse_expr >] -> App(App(Const Store, expr1), expr2)
    | [< 'Kwd "omega" >] -> parse_expr (lexer (Stream.of_string "while true do skip done"))
  
  let rec parse_full = parser (* produces a context and an expression *)
    | [< 'Kwd "-"; expr = parse_expr; 'Kwd "." >] -> ([], expr)
    | [< 'Ident x; 'Kwd ":"; typ = parse_type; full = parse_full >] -> ((x,typ)::(fst full), snd full)
    
  let parse_expr_str s = parse_expr (lexer (Stream.of_string s))
  let parse_type_str s = parse_type (lexer (Stream.of_string s))  
  let parse_full_str f = parse_full (lexer (Stream.of_string f))
  
end

(* SOME TOP LEVEL FUNCTIONS TO INTERFACE WITH THE ABOVE *)

(* Size of the datatype of natural numbers since we deal with
   the finitary sublanguage *)
let num_max = 8

(* Compiles a program string to a regexp *)
let regl s =
  let (c,t) = LangParse.parse_full_str s in
  Strategiser.den_sem num_max c t
  
(* Compiles a program string to an automaton *)  
let autol s =
  RegExp.compile_regexp (regl s)

(* Calculates the type of a program string *)
let typl s =
  let (c,t) = LangParse.parse_full_str s in
  Language.get_type c t

(* Reduces a closed program string to normal form *)
let redl s =
  let (c,t) = LangParse.parse_full_str s in
  if c = [] then Language.reduce t else failwith "Can only reduce closed term."

(* Gives an example play of a program string *)
let exl s =
  RegExp.example_word (regl s)
  
(* Returns if two program strings are equivalent *)
let eql s t =
  RegExp.equal (regl s) (regl t)

(* Returns if one program string refines another *)
let refl s t =
  RegExp.contains (regl s) (regl t)

(* Returns if two program strings are equivalent, and
   if not produces a seperating play that can be used
   to create a seperating context *)
let sepl s t =
  let (sr,tr) = (regl s, regl t) in
  let snott = RegExp.Both(sr, RegExp.Comp(RegExp.get_lang sr, tr)) in
  let tnots = RegExp.Both(tr, RegExp.Comp(RegExp.get_lang tr, sr)) in
  match RegExp.example_word snott with
    | Some x -> `FirstNotSecond x
    | None ->
          (match RegExp.example_word tnots with
             | Some x -> `SecondBigger x
             | None -> `Equal)

(* Determines whether a program string satisfies a
   property (represented as a regexp) *)
let satl s t =
  RegExp.contains (regl s) t

(* Example programs *)
let twops = "p : com; - new x in while < !x 2 do do set x = + !x 1 then p done ."
let twops2 = "p : com; - do p then p ."
let threeps = "p : com; - do p then do p then p ."

let no_sb = "p : com -> com; - new x in do @ p (set x = 1) then if (= !x 1) then omega else skip ."
let no_sb2 = "p : com -> com; - @ p omega ."

let local = "p : com; - p ."
let local2 = "p : com; - new x in p ."

let beta = " x : nat; - @ (fun y : nat -> + y 1) x ."
let beta2 = " x : nat; - + x 1 ."

let eta = " t : nat -> nat; - fun x : nat -> @ t x ."
let eta2 = " t : nat -> nat; - t ."

let even = " p : com -> com; - new x in do @ p (set x = + !x 2) then if even !x then omega else skip ."
let omega = " - omega ."

let four = " - @ (fun y : nat -> + y 3) 1 ."

(* Example properties / programs *)

open Strategiser
open Language
open RegExp

(* The property on the sequent p : com; x : var - _ : com saying x is written
   before p is called. *)
let xwbeforep =
  let context = [("p",Basic Com);("x",Basic VarN)] in
  let allmoves = Strategiser.all_moves num_max context (Basic Com) in
  let m = [Ctx("p",1,Run); Ctx("p",1,Done)] in
  let notm = List.filter (fun x -> not (List.mem x m)) allmoves in
  let restr xs = Star(alt (List.map (fun x -> Char x) xs)) in
  let writex = Sigma (num_max, (fun n -> Char (Ctx("x",1,Write (n-1))))) in
  seq [restr notm; writex; restr notm; Char (Ctx("p",1,Run)); restr allmoves]
  
let nosat_prop = "p : com; x : var; - p ."
let nosat_prop2 = "p : com; x : var; - do p then set x = 2 ."
let sat_prop = "p : com; x: var; - do set x = 1 then p ."