// Michael R. Hansen   22-11-2024

// Property-based testing for an expression language 
// substitution and let-elimination

#r "nuget: FsCheck";;

open FsCheck

type E =  | V of string
          | C of int
          | Let of string * E * E 
          | Add of E * E;;

type Env = Map<string,int>;;

// free: E -> Set<string>
let rec free = 
      function 
      | C _          -> Set.empty
      | V x          -> Set.singleton x
      | Add(e1,e2)   -> Set.union (free e1) (free e2)
      | Let(x,e1,e2) -> Set.union (free e1) (Set.remove x (free e2));;

//eval: E -> Env -> int
let rec eval e m = 
    match e with 
    | V x -> Map.find x m 
    | C n -> n
    | Let(x,e1,e) -> let v1 = eval e1 m 
                     eval e (Map.add x v1 m)
    | Add(e1,e2)  -> eval e1 m + eval e2  m;; 


// rename y y' e: renames all free occurrences of y in e to y'
let rec rename y y' = 
   function
   | V x when x=y -> V y'
   | Add(e1,e2)   -> Add(rename y y' e1, rename y y' e2)
   | Let(x,e1,e2) -> let e1' = rename y y' e1
                     if x=y then Let(x,e1',e2)
                     else Let(x,e1',rename y y' e2)
   | e            -> e;;

(* nextVar: string -> string *)
// nextVar x: generates a fresh (unused) variable from x. 
// makes use of a mutable cell n 
let nextVar = 
   let n = ref 0
   let f x = let s = x + "#" + string(n.Value) 
             (n.Value <- n.Value+1; s)
   f;; 

// subst: x e e': 
// The expression obtained from e' by substituting every free occurrence of x with e, where
// renaming of bound varibles in e' (so-called alpha-conversion) are performed to avoid clashes 
// where a free variable of e is captured by a let-binding in e'
let rec subst x e = 
   function 
   | V y when x=y          -> e
   | Add(e1,e2)            -> Add(subst x e e1, subst x e e2)
   | Let(y,e1,e2) when x=y -> Let(y, subst x e e1, e2)
   | Let(y,e1,e2)          -> let (y',e2') = if Set.contains y (free e) 
                                             then let y' = nextVar y
                                                  let e2' = rename y y' e2
                                                  (y',e2')
                                             else 
                                             (y,e2)      
                              Let(y', subst x e e1, subst x e e2')
   | e'                    -> e';;




// A fundamental property:
(* substOK1(x, e1, e2, m): 
  x:  variable used in a substitution
  e1: expression which should replace x
  e2: expression where the substitution is performed
  m:  env used for evaluation of expressions

  Substituting x with the meaning of e1 in e2 gives and expression that has the same meaning 
  as the expression obtained by substiting x with e1 in e2 
*)
type Subst = string * E * E * Env

let substOK1((x, e1, e2, m): Subst) = 
   let v1 = eval e1 m           
   let e2' = subst x (C v1) e2 
   let e2'' = subst x e1 e2     
   eval e2'' m = eval e2' m;;   

// This property assumes that m contains entries for the free variables occurring in the expressions. 
// Otherwise exceptions will be thrown by the eval function

// We need to generate well-formed values of type Subst when the property is checkted using PBT.


// Generators for Subst and EwithEnv:

// Generation of small strings, small small integers and environments for a list of variables

//  charsSeqGen: char -> chae -> Gen<char> seq 
let charsSeqGen c1 c2 = seq {for c in c1 .. c2 do
                               yield gen {return c} }

// mySmallStringGen: Gen<string>
let mySmallStringGen = Gen.map string (gen {return! Gen.oneof (charsSeqGen 'a' 'e')});;

// mySmallIntGen: Gen<int>
let mySmallIntGen = Gen.choose(0,9);;


// myEnvGen: 'a list -> Gen<Map<'a,int>
// genarator for environments to a list of variables vs
let myEnvGen vs = gen { let i   = List.length vs 
                        let! ns = Gen.listOfLength i mySmallIntGen 
                        return Map.ofList (List.zip vs ns)      };;


// myVarGen: 'a list -> Gen<'a>
// picks an element from vs
let myVarGen vs = 
    gen {let! i = Gen.choose(0, List.length vs - 1)
         return vs.[i] }

let myCGen    = Gen.map C mySmallIntGen;; 
let myVGen vs = Gen.map V (myVarGen vs);;

// myLeafGen: string list -> Gen<E>
// generator for leafs of type E on the basis of a list of variables
let myLeafGen vs = 
   if vs<>[] 
   then Gen.oneof [myCGen ; myVGen vs]
   else myCGen;;


// myEGen: string list -> Gen<E> 
// generates a value of type expression from a given list of variables vs
let myEGen vs =
   let rec myE vs n = 
      match n with 
      | 0 -> myLeafGen vs 
      | _ -> Gen.oneof [myLeafGen vs;
                        Gen.map2 (fun x y -> Add(x,y)) (myE vs (n/2)) (myE vs (n/2));
                        myLetGen vs n                                               ]
   and myLetGen vs n = gen {let! x  = mySmallStringGen 
                            let! e1  = myE vs (n/2) 
                            let! e = myE (x::vs) (n/2)
                            return Let(x,e1,e) }
   Gen.sized (myE vs);;
   
// Generator for well-formed values of type Subst
// - small comprehensible strings, integers and environments

let mySubstGen: Gen<Subst> =
   gen { let! i  = Gen.choose(1,6) 
         let! vs = Gen.listOfLength i mySmallStringGen 
         let! j  = Gen.choose(0,i-1)
         let x   = vs.[j]
         let! e1 = myEGen vs
         let! e2 = myEGen vs
         let! m  = myEnvGen (Set.toList( (Set.union (free e1) (free e2))))
         return (x,e1,e2,m) };

type EwithEnv = E * Env
let myEwithEnvGen: Gen<EwithEnv> = 
   gen { let! i = Gen.choose(1,6)                       
            // pick a number, call it i 
         
         let! vs = Gen.listOfLength i mySmallStringGen  
            // generate a variable list, call it vs
         
         let! e = myEGen vs                             
            // generate an expression from vs, call it e
         
         let! m = myEnvGen (Set.toList(free e))         
            // generate an environment for (free e), call it m
         
         return (e,m) }                                 

// Shrinkers are needed to get comprehensible counter-examples

// E -> E seq
let rec shrinkE e = 
   match e with 
   | V _ | C _    -> seq []
   | Add(e1,e2)   -> let se1 = shrinkE e1
                     let se2 = shrinkE e2
                     Seq.concat [seq [e1;e2]; se1; se2]
   | Let(x,e1,e2) -> let se1 = shrinkE e1
                     let se2 = shrinkE e2
                     let se2' = seq { for e2' in se2 do 
                                        if not (Set.contains x (free e2')) then
                                               yield e2' }
                     let se = seq { for e1' in se1 do 
                                        for e2' in se2 do
                                           yield Let(x,e1',e2')
                                         } 
                     Seq.concat [seq [e1]; se2'; se]

// Subst -> Subst seq
let mySubstShrinker(x,e1,e2,m) = 
   let e1s = shrinkE e1
   let e2s = shrinkE e2 
   seq { for e1' in e1s do 
             for e2' in e2s do 
                  yield (x, e1', e2', m)}

// EwithEnv -> EwithEnv seq
let myEwithEnvShrinker(e,m) = 
   seq { for e' in shrinkE e do
             yield (e',m)}

type MyGenerator =
  static member EwithEnv() =
      {new Arbitrary<EwithEnv>() with
          override x.Generator = myEwithEnvGen
          override x.Shrinker t = myEwithEnvShrinker t }
  static member SubstType() =
      {new Arbitrary<Subst>() with
          override x.Generator = mySubstGen
          override x.Shrinker t = mySubstShrinker t }
          ;;

Arb.register<MyGenerator>();;

let subst1 = Check.Quick substOK1;;


// Further properties and programs 


// elimLet: E -> E
let rec elimLet e = match e with
                    | Add(e1,e2)   -> Add(elimLet e1, elimLet e2)
                    | Let(x,e1,e2) -> elimLet (subst x e1 e2)
                    | e            -> e;;


let elimLetOK1((e,m): EwithEnv) = eval e m = eval (elimLet e) m;; 
let elimLet1 = Check.Quick elimLetOK1;;


// substEnv: Env -> E -> E 
let substEnv m e = Map.foldBack subst (Map.map (fun _ n -> C n) m) e

let substOK2((e,m): EwithEnv) = eval e m = eval (substEnv m e) Map.empty;;

let subst2 = Check.Quick substOK2;;


// all let-bindings are eliminated by elimLet

let rec nOfLet e = 
      match e with
      | Let(_,e1,e2) -> nOfLet e1 + nOfLet e2 + 1
      | Add(e1,e2)   -> nOfLet e1 + nOfLet e2
      | _            -> 0;;

let elimLetOK2 e = nOfLet(elimLet e) = 0

let elimLet2 = Check.Quick elimLetOK2;;


let rec toS e = match e with 
                  | V s -> s 
                  | C n -> string n
                  | Add(e1,e2) -> "(" + toS e1 + "+" + toS e2 + ")"
                  | Let(s,e1,e2) -> "let " + s + "=" + toS e1 + " in " + toS e2;;