// Link the F# PowerPack Linq library, for example:
#r @"C:\Program Files (x86)\FSharpPowerPack-1.9.9.9\bin\FSharp.PowerPack.Linq.dll"

open System
open System.Collections.Generic
open Microsoft.FSharp.Quotations
open Microsoft.FSharp.Quotations.Patterns
open Microsoft.FSharp.Linq.QuotationEvaluation


//--------------------------------------------------
// Expression traversal and mapping

let rec TraverseExpr f e =
    match f e with
    | Quote e -> Expr.Quote(TraverseExpr f e)
    | Let (n, v, e) -> Expr.Let(n, TraverseExpr f v, TraverseExpr f e)
    | LetRecursive (l, e) -> Expr.LetRecursive([for v,x in l -> v, TraverseExpr f x], TraverseExpr f e)
    | Lambda (v, e) -> Expr.Lambda(v, TraverseExpr f e)
    | Application (a, b) -> Expr.Application(TraverseExpr f a, TraverseExpr f b)
    | Call (None, m, l) -> Expr.Call(m, [for v in l -> TraverseExpr f v])
    | Call (Some e, m, l) -> Expr.Call(TraverseExpr f e, m, [for v in l -> TraverseExpr f v])
    | PropertyGet (None, p, l) -> Expr.PropertyGet(p, [for v in l -> TraverseExpr f v])
    | PropertyGet (Some e, p, l) -> Expr.PropertyGet(TraverseExpr f e, p, [for v in l -> TraverseExpr f v])
    | PropertySet (None, p, l, v) -> Expr.PropertySet(p, TraverseExpr f v, [for v in l -> TraverseExpr f v])
    | PropertySet (Some e, p, l, v) -> Expr.PropertySet(TraverseExpr f e, p, TraverseExpr f v, [for v in l -> TraverseExpr f v])
    | FieldGet (Some e, d) -> Expr.FieldGet(TraverseExpr f e, d)
    | FieldSet (None, d, v) -> Expr.FieldSet(d, TraverseExpr f v)
    | FieldSet (Some e, d, v) -> Expr.FieldSet(TraverseExpr f e, d, TraverseExpr f v)
    | NewTuple l -> Expr.NewTuple [for v in l -> TraverseExpr f v]
    | NewRecord (t, l) -> Expr.NewRecord(t, [for v in l -> TraverseExpr f v])
    | NewArray (t, l) -> Expr.NewArray(t, [for v in l -> TraverseExpr f v])
    | NewUnionCase (u, l) -> Expr.NewUnionCase(u, [for v in l -> TraverseExpr f v])
    | NewObject (c, l) -> Expr.NewObject(c, [for v in l -> TraverseExpr f v])
    | Coerce (e, t) -> Expr.Coerce(TraverseExpr f e, t)
    | TypeTest (e, t) -> Expr.TypeTest(TraverseExpr f e, t)
    | UnionCaseTest (e, u) -> Expr.UnionCaseTest(TraverseExpr f e, u)
    | TupleGet (e, i) -> Expr.TupleGet(TraverseExpr f e, i)
    | IfThenElse (a, b, c) -> Expr.IfThenElse(TraverseExpr f a, TraverseExpr f b, TraverseExpr f c)
    | Sequential (a, b) -> Expr.Sequential(TraverseExpr f a, TraverseExpr f b)
    | ForIntegerRangeLoop (v, a, b, c) -> Expr.ForIntegerRangeLoop(v, TraverseExpr f a, TraverseExpr f b, TraverseExpr f c)
    | WhileLoop (a, b) -> Expr.WhileLoop(TraverseExpr f a, TraverseExpr f b)
    | TryWith (a, v, b, w, c) -> Expr.TryWith(TraverseExpr f a, v, TraverseExpr f b, w, TraverseExpr f c)
    | TryFinally (a, b) -> Expr.TryFinally(TraverseExpr f a, TraverseExpr f b)
    | AddressOf e -> Expr.AddressOf(TraverseExpr f e)
    | AddressSet (a, b) -> Expr.AddressSet(TraverseExpr f a, TraverseExpr f b)
    | VarSet (v, e) -> Expr.VarSet(v, TraverseExpr f e)
    | NewDelegate (t, l, e) -> Expr.NewDelegate(t, l, TraverseExpr f e)
//  | Var v -> 
//  | Value (a, t) ->
//  | DefaultValue t -> 
    | e -> e


//--------------------------------------------------
// The expression transformers for rename and reduce

module internal TransformExpr =
    let vs = Dictionary<string, int>()
    let splitname (s: string) =
        let mutable c = ref 0
        match s.Split [|'_'|] with
        | [|s|] -> s,0
        | [|s;n|] when Int32.TryParse(n, c) -> s,!c
        | _ -> failwith <| "Variable " + s + " is not properly named. Do not use _ in a var name."

    let GetNew name typ =
        let s,_ = splitname name
        if not (vs.ContainsKey s) then vs.Add(s,1) else vs.[s] <- vs.[s] + 1
        Var.Global(s + "_" + (string vs.[s]), typ)

    let unique (v: Var) =
        let s,n = splitname v.Name
        if not (vs.ContainsKey s) || n=0 then GetNew s v.Type
        else if Var.Global(v.Name, v.Type) <> v then GetNew s v.Type
        else v
    let fix v1 v2 (e: Expr) =
        if v1 = v2 then e
        else
            let nv = Some (Expr.Var v2)
            e.Substitute (fun v -> if v=v1 then nv else None)

    let MakeUnique = function
        | Var v -> Expr.Var (unique v)
        | Lambda (n, e) -> let n' = unique n in Expr.Lambda(n', fix n n' e)
        | Let (n, v, e) -> let n' = unique n in Expr.Let(n', v, fix n n' e)
        | ForIntegerRangeLoop (n, a, b, e) -> let n' = unique n in Expr.ForIntegerRangeLoop(n, a, b, fix n n' e)
        | x -> x

    let MakeSplice = function
        | Application (Call(None, m, [Value(p,_)]), v)
            when (m.Name = "SpliceExpression") &&
                 (m.DeclaringType.ToString() = "Microsoft.FSharp.Core.ExtraTopLevelOperators") &&
                 (p :? Expr) ->
            Expr.Application(p :?> Expr, v)
        | x -> x


// Get a new unique typed variable
let Variable<'v> name = TransformExpr.GetNew name typeof<'v>

// Unique variable rename an expression
let Rename (x: Expr<'a>) : Expr<'a> = TraverseExpr TransformExpr.MakeUnique x |> Expr.Cast
// Code reduce an expression
let Reduce (x: Expr<'a>) : Expr<'a> = TraverseExpr TransformExpr.MakeSplice x |> Expr.Cast
// Perform all expression transformations
let Transform (x: Expr<'a>) : Expr<'a> = TraverseExpr (TransformExpr.MakeUnique >> TransformExpr.MakeSplice) x |> Expr.Cast


//--------------------------------------------------
// The quotations pretty printer

module internal QuotationsPrinter =
    let opNameTable = 
       [("[]", "op_Nil"); ("::", "op_ColonColon"); ("+", "op_Addition"); (":=", "op_ColonEquals"); ("ref", "Ref");
        ("~%%", "op_SpliceUntyped"); ("~++", "op_Increment"); ("~--", "op_Decrement"); ("-", "op_Subtraction");
        ("*", "op_Multiply"); ("**", "op_Exponentiation"); ("/", "op_Division"); ("@", "op_Append");
        ("^", "op_Concatenate"); ("%", "op_Modulus"); ("&&&", "op_BitwiseAnd"); ("|||", "op_BitwiseOr");
        ("^^^", "op_ExclusiveOr"); ("<<<", "op_LeftShift"); ("~~~", "op_LogicalNot"); (">>>", "op_RightShift");
        ("~+", "op_UnaryPlus"); ("~-", "op_UnaryNegation"); ("~&", "op_AddressOf"); ("~&&", "op_IntegerAddressOf");
        ("&&", "op_BooleanAnd"); ("||", "op_BooleanOr"); ("<=", "op_LessThanOrEqual"); ("=","op_Equality");
        ("<>","op_Inequality"); (">=", "op_GreaterThanOrEqual"); ("<", "op_LessThan"); (">", "op_GreaterThan");
        ("|>", "op_PipeRight"); ("||>", "op_PipeRight2"); ("|||>", "op_PipeRight3"); ("<|", "op_PipeLeft");
        ("<||", "op_PipeLeft2"); ("<|||", "op_PipeLeft3"); ("!", "op_Dereference"); (">>", "op_ComposeRight");
        ("<<", "op_ComposeLeft"); ("+=", "op_AdditionAssignment"); ("-=", "op_SubtractionAssignment");
        ("*=", "op_MultiplyAssignment"); ("/=", "op_DivisionAssignment"); ("..", "op_Range"); ("?", "op_Dynamic")]

    let IsOpName = 
        let t = new Dictionary<_,_>()
        for (s,n) in opNameTable do t.Add(n,s)
        fun (n:string) -> if t.ContainsKey n then t.[n] else n

    let rec PrintQuote (x: Expr) =
        let inst = function Some e -> "(" + PrintQuote e + ")." | None -> ""
        let args l b = if l = [] && b then "" else "(" + String.concat ", " (List.map PrintQuote l) + ")"
        match x with
        | Var v -> v.Name
        | Value (null, _) -> "()"
        | Value (a, _) when (a :? string) -> "\"" + a.ToString() + "\""
        | Value (a, _) -> a.ToString()
        | Quote e -> "<@ " + PrintQuote e + " @>"
        | Let (n, v, e) -> "let " + n.Name + " = " + PrintQuote v + " in " + PrintQuote e
        | LetRecursive (l, e) -> "let rec " + String.concat " and " (List.map (fun (n:Var,v:Expr) -> n.Name + " = " + PrintQuote v) l) + " in " + PrintQuote e
        | Lambda (v, e) -> "fun " + v.Name + " -> " + PrintQuote e
        | Application (a, b) -> "(" + PrintQuote a + ") (" + PrintQuote b + ")"
        | Call (None, m, [p1]) when m.Name <> IsOpName m.Name -> IsOpName m.Name + "(" + PrintQuote p1 + ")"
        | Call (None, m, [p1;p2]) when m.Name <> IsOpName m.Name -> "(" + PrintQuote p1 + ")" + IsOpName m.Name + "(" + PrintQuote p2 + ")"
        | Call (i, m, l) -> (inst i) + m.ReflectedType.FullName + "." + m.Name + args l false
        | PropertyGet (i, p, l) -> inst i + p.ReflectedType.FullName + "." + p.Name + args l true
        | PropertySet (i, p, l, v) -> inst i + p.ReflectedType.FullName + "." + p.Name + args l true + " <- " + PrintQuote v
        | FieldGet (i, f) -> inst i + f.ReflectedType.FullName + "." + f.Name
        | FieldSet (i, f, v) -> inst i + f.ReflectedType.FullName + "." + f.Name + " <- " + PrintQuote v
        | NewTuple l -> args l false
        | NewRecord (t, l) -> "({" + String.concat "; " (List.map PrintQuote l) + "}:" + t.Name + ")"
        | NewArray (_, l) -> "[|" + String.concat "; " (List.map PrintQuote l) + "|]"
        | NewUnionCase (u, l) -> u.Name + args l true
        | NewObject (c, l) -> "new " + c.DeclaringType.FullName + args l false
        | Coerce (e, t) -> "(" + PrintQuote e + " :> " + t.Name + ")"
        | TypeTest (e, t) -> "(" + PrintQuote e + " :? " + t.Name + ")"
        | UnionCaseTest (e, u) -> "(" + PrintQuote e + " = " + u.Name + ")"
        | TupleGet (e, i) -> "TupleGet(" + PrintQuote e + ", " + i.ToString() + ")"
        | IfThenElse (a, b, c) -> "if " + PrintQuote a + " then " + PrintQuote b + " else " + PrintQuote c
        | Sequential (a, b) -> "(" + PrintQuote a + " ; " + PrintQuote b + ")"
        | ForIntegerRangeLoop (v, a, b, c) -> "for " + v.Name + " = " + PrintQuote a + " to " + PrintQuote b + " do " + PrintQuote c + " done"
        | WhileLoop (a, b) -> "while " + PrintQuote a + " do " + PrintQuote b + " done"
        | TryWith (a, _, _, v, b) -> "try " + PrintQuote a + " with " + v.Name + " -> " + PrintQuote b
        | TryFinally (a, b) -> "try " + PrintQuote a + " finally " + PrintQuote b
        | AddressOf e -> PrintQuote e
        | AddressSet (a, b) -> PrintQuote a + " <- " + PrintQuote b
        | VarSet (v, e) -> "VarSet(" + v.Name + ", " + PrintQuote e + ")"
        | DefaultValue t -> "(default " + t.Name + ")"
        | NewDelegate (t, l, e) -> x.ToString()
        | _ -> "(* Can't pretty print this: *) (" + x.ToString() + ")"


// Pretty-print a quotation
let GetQuotedCode q = "<@ " + QuotationsPrinter.PrintQuote q + " @>"

// To add the pretty printer as default quotation printer in F# interactive:
// fsi.AddPrinter GetQuotedCode


//--------------------------------------------------
// Combinators

// Basic combinators and values for quoted code
let Lift x = <@ x @>
let Eval (x: Expr<'a>) = x.Eval()
let Unit = <@ () @>

// let statement
let Let name (a: Expr<'v>) (f: Expr<'v> -> Expr<'e>) : Expr<'e> = 
    let v = Variable<'v> name
    Expr.Let(v, a, f (v |> Expr.Var |> Expr.Cast)) |> Expr.Cast

// Function abstraction
// It is the user's responsibility to call this when they have a lambda function code. Multiple parameter lambdas are OK
let Lambda name (f: Expr<'a> -> Expr<'b>) : Expr<'a->'b> =
    let v = Variable<'a> name
    Expr.Lambda(v, f (v |> Expr.Var |> Expr.Cast)) |> Expr.Cast
let LambdaL (f: Expr<'a->'b>) : Expr<'a->'b> = Transform f

// Function application
let Apply (f: Expr<'a->'b>) (x: Expr<'a>) : Expr<'b> = Expr.Application(Transform f, Transform x) |> Expr.Cast

// Sequencing
let Sequence (a: Expr<'b>) (b: Expr<'a>) : Expr<'a> = Expr.Sequential(Transform a, Transform b) |> Expr.Cast

// Other control structures
let While c b = let c', b' = Transform c, Transform b in <@ while (%c') do (%b') done @>
let If cd th el = let c, t, e = Transform cd, Transform th, Transform el in <@ if (%c) then (%t) else (%e) @>
let TryWith b w = let b', w' = Transform b, Transform w in <@ try (%b') with _ -> (%w') @>
let TryFinally b f = let b', f' = Transform b, Transform(f()) in <@ try (%b') finally (%f') @>

// Ease of use operators
let Ignore a = let a' = Transform a in <@ ignore (%a') @>
let Ref c = <@ ref (%c) @>
let (!@) c = <@ !(%c) @>
let (=@) c v = let v' = Transform v in <@ (%c) := (%v') @>
let SomeL v = <@ Some (%v) @>
let OptionL = function Some v -> SomeL v | _ -> <@ None @>


//--------------------------------------------------
// Code Generation workflow

type CGBuilder() =
    member cg.Delay v = v()
    member cg.Bind (v,f) = Let "t" v f
    member cg.Return v = <@ v @>
    member cg.ReturnFrom v = v
    member cg.YieldFrom v = v
    member cg.Combine (a,b) = Sequence a b
    member cg.TryFinally (b,f) = TryFinally b f
    member cg.While (c,e) = While (c()) e
    member cg.Zero () = <@ () @>
let codegen = CGBuilder()


//--------------------------------------------------
// Code Generation workflow with state and explicit CPS notation

// state + continuation monad
type StateCPSMonad<'state, 'value, 'answer> = 'state -> ('state -> 'value -> 'answer) -> 'answer
type SCMonadBuilder() =
    member x.Return a = fun s k -> k s a
    member x.Bind (m, f) = fun s k -> m s (fun s' k' -> f k' s' k)
let StateCPS = SCMonadBuilder()
let k0 _ v = v
let runM m s0 = m s0 k0
let fetch s k = k s s
let store v _ k = k v ()

// Code generation StateCPS monad. State is a list, values and answers are quoted expressions
// This overrides everything in the StateCPS monad for type enforcement on function signatures
type CodeGen<'s, 'v, 'w> = StateCPSMonad<seq<'s>, Expr<'v>, Expr<'w>>
type CGMonadBuilder() =
    member x.Return a : CodeGen<_,_,_> =
        fun s k -> k s a
    member x.Bind (m: CodeGen<_,_,_>, f: _ -> CodeGen<_,_,_>) : CodeGen<_,_,_> =
        fun s k -> m s (fun s' k' -> f k' s' k)
let mcodegen = CGMonadBuilder()

// Generating code from mcodegen monad, and reseting
let ResetM _ = id
let GenerateM (m: CodeGen<_,_,_>) = m [] ResetM

// State morphisms
let GetStateM : CodeGen<_,_,_> = fun s k -> k s <@ s @>
let SetStateM v : CodeGen<_,_,_> = fun _ k -> k v <@ () @>

// StateCPS Monadic code combinators
let UnitM: CodeGen<_,_,_> = fun s k -> k s Unit
let ReturnM v = fun s k -> k s v
let LetM name a = fun s k -> Let name a (k s)
let ApplyM f x = fun s k -> k s (Apply (f s ResetM) (x s ResetM))
let SequenceM a b = fun s k -> k s (Sequence (a s ResetM) (b s ResetM))
let Prepend a b = fun s k -> Sequence a (k s b)
let PrependM a b = fun s k -> Sequence (a s ResetM) (k s b)
let WhileM c b = fun s k -> k s (While (c s ResetM) (b s ResetM))
let IterM it fn v = fun s k -> k s <| Apply (Apply (Transform it) <| Lambda "a" (fun a -> fn a s ResetM)) (Transform v)
let IfM cd th el  = fun s k -> k s (If cd (th s ResetM) (el s ResetM))
let TryWithM b w = fun s k -> k s (TryWith (b s ResetM) (w s ResetM))
let TryFinallyM b f = fun s k -> k s (TryFinally (b s ResetM) (f s ResetM))
let IgnoreM a = fun s k -> k s (Ignore (a s ResetM))