safeil.fs

module SafeIL

open System
open System.Reflection.Emit
open System.Runtime.CompilerServices
open System.Runtime.InteropServices
open System.Reflection
open Microsoft.FSharp.Quotations
open Microsoft.FSharp.Quotations.Patterns
open Microsoft.FSharp.Quotations.DerivedPatterns

// Phantom types to track stack types and parameter types:
module Stack =
    type Empty = private | Empty
    type Cons<'t, 'ts> = private | Cons

module Args =
    type None = private | Empty
    type Then<'t, 'ts> = private | Then

// TODO: need VoidCompletedFunction as well, *shakes fist at CLR*
type CompletedFunction<'args, 'ret> = CompletedIL of Type[] * (ILGenerator -> unit)
type Function<'args, 'stack> = IL of Type[] * (ILGenerator -> unit)

let inline private emit f (IL (ts, oldF)) = IL (ts, (fun g -> oldF g; f g))

[<AutoOpen>]
[<Extension>]
type Ops =

    static member method () : Function<Args.None, Stack.Empty> = IL (Type.EmptyTypes, fun _ -> ())
    static member method<'a> () : Function<Args.Then<'a, Args.None>, Stack.Empty> = IL ([|typeof<'a>|], fun _ -> ())
    static member method<'a, 'b> () : Function<Args.Then<'a, Args.Then<'b, Args.None>>, Stack.Empty> = IL ([|typeof<'a>; typeof<'b>|], fun _ -> ())

    static member push (x : int) : Function<'args, 'stack> -> Function<'args, Stack.Cons<int, 'stack>> =
        emit (fun g -> g.Emit(OpCodes.Ldc_I4, x))

    static member push (x : int64) : Function<'args, 'stack> -> Function<'args, Stack.Cons<int, 'stack>> =
        emit (fun g -> g.Emit(OpCodes.Ldc_I8, x))

    static member push (s : string) : Function<'args, 'stack> -> Function<'args, Stack.Cons<string, 'stack>> =
        emit (fun g -> g.Emit(OpCodes.Ldstr, s))

    static member pop (f : Function<'args, Stack.Cons<'a, 'rest>>) : Function<'args, 'rest> =
        emit (fun g -> g.Emit OpCodes.Pop) f

    static member dup (f : Function<'args, Stack.Cons<'a, 'rest>>) : Function<'args, Stack.Cons<'a, Stack.Cons<'a, 'rest>>> =
        emit (fun g -> g.Emit OpCodes.Dup) f

    static member add (f : Function<'args, Stack.Cons<int, Stack.Cons<int, 'rest>>>) : Function<'args, Stack.Cons<int, 'rest>> =
        emit (fun g -> g.Emit OpCodes.Add) f

    static member loadArg_0 (f : Function<Args.Then<'a, 'args>, 'stack>) : Function<Args.Then<'a, 'args>, Stack.Cons<'a, 'stack>> =
        emit (fun g -> g.Emit OpCodes.Ldarg_0) f

    static member loadArg_1 (f : Function<Args.Then<'a, Args.Then<'b, 'args>>, 'stack>) : Function<Args.Then<'a, Args.Then<'b, 'args>>, Stack.Cons<'b, 'stack>> =
        emit (fun g -> g.Emit OpCodes.Ldarg_1) f

    static member loadArg_2 (f : Function<Args.Then<'a, Args.Then<'b, Args.Then<'c, 'args>>>, 'stack>) : Function<Args.Then<'a, Args.Then<'b, Args.Then<'c, 'args>>>, Stack.Cons<'c, 'stack>> =
        emit (fun g -> g.Emit OpCodes.Ldarg_2) f

    static member loadArg_3 (f : Function<Args.Then<'a, Args.Then<'b, Args.Then<'c, Args.Then<'d, 'args>>>>, 'stack>) : Function<Args.Then<'a, Args.Then<'b, Args.Then<'c, Args.Then<'d, 'args>>>>, Stack.Cons<'d, 'stack>> =
        emit (fun g -> g.Emit OpCodes.Ldarg_3) f

    static member ret (IL (ts, insns) : Function<'args, Stack.Cons<'r, Stack.Empty>>) : CompletedFunction<'args, 'r> =
        CompletedIL (ts, fun g -> insns g; g.Emit OpCodes.Ret)
    
    static member ret (IL (ts, insns) : Function<'args, Stack.Empty>) : CompletedFunction<'args, unit> =
        CompletedIL (ts, fun g -> insns g; g.Emit OpCodes.Ret)

    static member ifZero (z : Function<'args, 'stack> -> Function<'args, 'newStack>) (nz : Function<'args, 'stack> -> Function<'args, 'newStack>) (IL (ts, insns) as f : Function<'args, Stack.Cons<int, 'stack>>) : Function<'args, 'newStack> =
        emit (fun g ->
            let dummyFun = IL (ts, (fun _ -> ())) // new starter fun
            let whenZero = g.DefineLabel()
            g.Emit(OpCodes.Brfalse, whenZero)
            let (IL (_, nzI)) = nz dummyFun in nzI g
            let phiLabel = g.DefineLabel()
            g.Emit(OpCodes.Br, phiLabel)
            g.MarkLabel whenZero
            let (IL (_, zI)) = z dummyFun in zI g
            g.MarkLabel phiLabel
        ) f

    static member callFunc (expr : Expr<unit -> 'a>) : Function<'args, 'stack> -> Function<'args, Stack.Cons<'a, 'stack>> =
        match expr with
        | Lambda (v, Call (None, mi, args)) ->
            emit (fun g -> g.Emit(OpCodes.Call, mi))
        | _ -> raise <| ArgumentException "Argument to call must be a lamdba"

    static member callAction (expr : Expr<'a -> unit>) : Function<'args, Stack.Cons<'a, 'stack>> -> Function<'args, 'stack> =
        match expr with
        | Lambda (v, Call (None, mi, args)) ->
            emit (fun g -> g.Emit(OpCodes.Call, mi))
        | _ -> raise <| ArgumentException "Argument to call must be a lamdba"

    static member callAction (expr : Expr<'a -> 'b -> unit>) : Function<'args, Stack.Cons<'b, Stack.Cons<'a, 'stack>>> -> Function<'args, 'stack> =
        match expr with
        | Lambda (v, Call (None, mi, args)) ->
            emit (fun g -> g.Emit(OpCodes.Call, mi))
        | _ -> raise <| ArgumentException "Argument to call must be a lamdba"

    static member call (expr : Expr<'a -> 'b>) : Function<'args, Stack.Cons<'a, 'stack>> -> Function<'args, Stack.Cons<'b, 'stack>> =
        if typeof<unit> = typeof<'a>
        then raise <| ArgumentException "Use callFunc"

        if typeof<unit> = typeof<'b>
        then raise <| ArgumentException "Use callAction"

        match expr with
        | Lambda (v, Call (None, mi, args)) ->
            emit (fun g -> g.Emit(OpCodes.Call, mi))
        | _ -> raise <| ArgumentException "Argument to call must be a lamdba"

    static member call (expr : Expr<'a -> 'b -> 'c>) : Function<'args, Stack.Cons<'b, Stack.Cons<'a, 'stack>>> -> Function<'args, Stack.Cons<'c, 'stack>> =
        if typeof<unit> = typeof<'a> || typeof<unit> = typeof<'b>
        then raise <| ArgumentException "Unit arguments?!"

        if typeof<unit> = typeof<'b>
        then raise <| ArgumentException "Use callAction"

        match expr with
        | Lambda (v, Call (None, mi, args)) ->
            emit (fun g -> g.Emit(OpCodes.Call, mi))
        | _ -> raise <| ArgumentException "Argument to call must be a lamdba"

let defineStaticMethod (tb : TypeBuilder) (name : string) (CompletedIL (ts, insns) : CompletedFunction<'Args, 'r>) =
    let resultType = if typeof<'r> = typeof<unit> then typeof<Void> else typeof<'r>
    let m = tb.DefineMethod(name, MethodAttributes.Static ||| MethodAttributes.Public, resultType, ts) // TODO: visibility
    insns <| m.GetILGenerator()
    m

let defineInstanceMethod (tb : TypeBuilder) (name : string) (CompletedIL (ts, insns) : CompletedFunction<Args.Then<'This, 'Args>, 'r>) =
    if typeof<'This> <> (tb :> Type)
    then raise <| ArgumentException("Wrong type for 'this'")

    let m = tb.DefineMethod(name, MethodAttributes.Public, typeof<'r>, ts) // TODO: visibility
    insns <| m.GetILGenerator()
    m

[<EntryPoint>]
let main argv = 
    let dynAsm = AssemblyBuilder.DefineDynamicAssembly(AssemblyName "A", AssemblyBuilderAccess.RunAndSave)
    let dynMod = dynAsm.DefineDynamicModule("M")
    let t = dynMod.DefineType("T")

    // if the sum is zero, prints -100, otherwise prints 100, then returns the sum
    let adder : CompletedFunction<Args.None, int> = 
        Ops.method()
        |> Ops.push "X: "
        |> Ops.callAction <@ (fun (s : String) -> Console.WriteLine(s)) @>
        |> Ops.callFunc <@ fun () -> Console.ReadLine() @>
        |> Ops.call <@ fun (s : String) -> Int32.Parse(s) @>
        |> Ops.push "Y: "
        |> Ops.callAction <@ (fun (s : String) -> Console.WriteLine(s)) @>
        |> Ops.callFunc <@ fun () -> Console.ReadLine() @>
        |> Ops.call <@ fun (s : String) -> Int32.Parse(s) @>
        |> Ops.add
        |> Ops.dup
        |> Ops.ifZero
            (Ops.push "Sum is zero!")
            (Ops.push "Sum is not zero!")
        |> Ops.callAction <@ fun (x : String) -> Console.WriteLine(x) @>
        |> Ops.ret

    let mb = defineStaticMethod t "Adder" adder
    let dt = t.CreateType()

    let method = dt.GetMethod("Adder")
    let result = method.Invoke(null, [| |]) :?> int
    printfn "%d" result
    
    0