// The ring example for generating power functions

open fscodegen


type Ring<'t> =
    abstract zero: 't
    abstract one: option<'t>
    abstract add: 't -> 't -> 't
    abstract neg: 't -> 't
    abstract sub: option<'t -> 't -> 't>
    abstract mul: 't -> 't -> 't
    abstract inv: option<'t -> 't>
    abstract div: option<'t -> 't -> 't>


let mksub (r: Ring<_>) =
    match r.sub with
    | None ->
        Some (fun a b -> r.add a (r.neg b))
    | _ -> r.sub
let mkdiv (r: Ring<_>) =
    match r.div, r.inv with
    | None, Some inv ->
        Some (fun a b -> r.mul a (inv b))
    | _, _ -> r.div


let Ints = {
    new Ring<_> with
    member i.zero = 0
    member i.one = Some 1
    member i.add a b = a + b
    member i.neg a = -a
    member i.sub = Some (fun a b -> a - b)
    member i.mul a b = a * b
    member i.inv = None
    member i.div = None
}

let LiftedInts = {
    new Ring<_> with
    member i.zero = <@ 0 @>
    member i.one = Some <@ 1 @>
    member i.add a b = <@ (%a) + (%b) @>
    member i.neg a = <@ -(%a) @>
    member i.sub = None
    member i.mul a b = <@ (%a) * (%b) @>
    member i.inv = None
    member i.div = None
}

let LiftedReals = {
    new Ring<_> with
    member i.zero = <@ 0. @>
    member i.one = Some <@ 1. @>
    member i.add a b = <@ (%a) + (%b) @>
    member i.neg a = <@ -(%a) @>
    member i.sub = Some (fun a b -> <@ (%a) - (%b) @>)
    member i.mul a b = <@ (%a) * (%b) @>
    member i.inv = Some (fun a -> <@ 1. / (%a) @>)
    member i.div = Some (fun a b -> <@ (%a) / (%b) @>)
}


let rec Power (r: Ring<_>) n b =
    match n with
    | 1 -> b
    | n when n%2 = 1 -> codegen {
        let! b' = r.mul b b
        yield! r.mul (Power r (n/2) b') b }
    | n -> codegen {
        let! b' = r.mul b b
        yield! Power r (n/2) b' }

let Power' r n b =
    match n with
    | n when n>0 -> Power r n b |> SomeL
    | 0 -> OptionL r.one
    | n ->
        match r.inv with
        | Some inv -> Power r (-n) b |> inv |> SomeL
        | _ -> <@ None @>


let MakePower r n = Power' r n |> Lambda "v"