open Prelude
open StateCPSMonad

type 'b abstract = unit -> 'b

(* types shared between code, direct and coderep *)
type ('pc,'p) cmonad_constraint = unit
    constraint
        'p = <state : 's list; answer : 'w abstract>
    constraint
        'pc = <answer : 'w; state : 's; ..>

type ('pc,'v) cmonad = ('p, 'v abstract) monad
    constraint _ = ('pc,'p) cmonad_constraint

type ('p, 'a) same01 =
    Same of ('a list abstract -> ('p, 'a list) cmonad)
  | Different of ('p, 'a list) cmonad * ('a abstract -> ('p,'a list) cmonad)

(* This one is very special! *)
let retN a = fun s k -> fun () -> let t = a () in k s (fun () -> t) ()

(* Naming conventions: 
   Functions that end in M take monads as arguments and partially
   run them -- in a reset fashion
   Functions that end in L always return abstract values from plain
   value -- the L is stands for ``lifted'' *)

(* sequencing *)
(* Note: the difference between `seq' and `seqM' is quite akin
   to the difference between call-by-value and call-by-name.
   Also note that `seq' can be expressed in terms of `seqM',
   but not the other way around 
*)

let seq a b = fun () -> begin a (); b () end
let seqL a b = fun s k -> k s (seq a b)

let seqM a b = fun s k -> k s (fun () ->  begin (a s k0) () ; (b s k0) () end )


let seqMS a b = liftM2 seq a b

let optSeq a = function
    | Some c -> seq a c
    | None   -> a

let optSeqM a = function
    | Some c -> seqM a c
    | None   -> a

(* conditional *)
(* Note the implicit `reset'.  Also note how the condition does not
   involve a `reset' *)
(* Note the implicit `reset' *)
let ifM test th el = fun s k ->
  k s (fun () ->  if test () then th s k0 () else el s k0 () )

let ifL test th el = fun () -> if test () then th () else el ()

let rshiftM cf = fun s k -> k s (cf s)

let whenM test th  = rshiftM (fun s -> 
  fun () ->  if test () then th s k0 () else () )

(* loops actually bind a value *)
let loopM low high body = function
  | UP -> (fun s k -> 
    k s (fun () ->  for j = low () to high () do body (fun () -> j) s k0 () done
))
  | DOWN -> (fun s k -> 
    k s (fun () ->  for j = low () downto high () do body (fun () -> j) s k0 ()
done ))

(* while ``loops'' do not naturally bind a value *)
let whileM cond body = fun s k -> 
    k s (fun () ->  while cond () do body s k0 () done )

let matchOpt x consequence1 consequence2 = fun () ->
  match x () with
    | None -> consequence1 ()
    | Some n -> (consequence2 (fun () -> n)) ()

(* match for Some/None *)
let matchM x som non = fun s k ->
    k s (fun () ->  match x () with
           | Some i -> som (fun () -> i) s k0 ()
           | None   -> non s k0 () )


(* generic loop *)
let genrecloop gen rtarg = fun s k -> k s
    (fun () -> let rec loop j = gen (fun () -> loop) (fun () -> j) s k0 () in
    loop (rtarg ()))

(* another non-trivial morphism: generate a bit of code *)
(* let codegen v cf = fun s k -> cf (k s v) *)

(* set up some very general algebra that can be reused though not
   so much in the current code (yet?) 
type ('a,'b) abstract = Ground of 'b | Code of ('a, 'b) code
let concretize = function
    | Ground x -> .<x>.
    | Code x   -> x
*)

(* We now define a lot of infrastructure.  We will be quite
   thorough, and define base abstraction and lifted abstractions,
   This will involve a lot of boilerplate code, which unfortunately 
   cannot be so easily automated in MetaOCaml -- that would require 
   introspection.  It could be done in camlp4, but that seems too 
   much as well.  This 'base' could be elided, but when we decide
   to make more use of Abstract Interpretation, we'll regret it,
   so do it now. *)

let lift x = fun () -> x

let liftRef x = fun () -> ref (x ())
let liftGet x = fun () -> ! (x ())

(* Need to use `fun' explictly to avoid monomorphising *)
let unitL = fun s k -> k s (fun () -> ())

module Pair = struct
(* To be able to deconstruct pairs in monadic code:
   perform (a,b) <-- ret (unpair pv) *)
  let unpair x = ((fun () -> fst (x ())), fun () -> snd (x ()))
  let fst x = (fun () -> fst (x ()))
  let fstM x = fun s k -> fun () -> let (a,_) = x () in k s (fun () -> a) ()
  let snd x = (fun () -> snd (x ()))
end

let liftPPair x = ((fun () -> fst (fst (x ()))) , (fun () -> snd (fst (x ()))),
fun () -> snd (x ()) )

let liftSome x = fun () -> Some (x ())

(* logic code combinators - plain and monadic *)
module Logic = struct
  let trueL         = fun () -> true
  let falseL        = fun () -> false
  let notL a        = fun () -> not (a ())
  let equalL a b    = fun () -> (a ()) = (b ())
  let notequalL a b = fun () -> (a ()) <> (b ())
  let andL a b     = fun () -> (a ()) && (b ())
end

module Size = struct
  let size n = fun () -> n
end

(* operations on code indices *)
module Idx = struct
  let zero = fun () -> 0
  let one = fun () -> 1
  let minusone = fun () -> -1
  let succ a = fun () -> (a ()) + 1
  let pred a = fun () -> (a ()) - 1
  let less a b = fun () -> (a ()) < (b ())
  let ltoe a b = fun () -> (a ()) <= (b ())
  let uminus a = fun () -> - (a ())
  let add a b = fun () -> (a ()) + (b ())
  let sub a b = fun () -> (a ()) - (b ())

  (* need explicit fun to avoid monomorphising *)
  let minusoneL = fun s k -> k s minusone
end

(* Maybe code generator *)
module Maybe = struct
  let just x = fun () -> Some (x ())
  let none   = fun () -> None
end

let applyMaybe g x = match g with
    | Some f -> f x
    | None   -> x

(* monadic tuples *)
module Tuple = struct
  let tup2 a b = fun () -> ( (a ()), (b ()) )
  let tup3 a b c = fun () -> ( (a ()), (b ()), (c ()) )
  let tup4 a b c d = fun () -> ( (a ()), (b ()), (c ()), (d ()) )
end

(* List ops *)
module CList = struct
  let nil = fun () -> []
  let cons a b = fun () -> (a ()) :: (b ())
  let length l = fun () -> List.length (l ())
  let nth l i = fun () -> List.nth (l ()) (i ())
  let hd l = fun () -> List.hd (l ())
  let tl l = fun () -> List.tl (l ())
  let append a b = fun () -> List.append (a ()) (b ())
  let matchList l empty nonempty = fun () ->
    match l () with
      | []   -> empty ()
      | h::t -> nonempty (fun () -> h) (fun () -> t) ()
  let matchListM l empty nonempty = fun s k -> (fun () ->
    match l () with
      | []   -> empty s k ()
      | h::t -> nonempty (fun () -> h) (fun () -> t) s k ())
  let matchListTwoM l emptyorone twoplus = fun s k ->
    match emptyorone with
    | Same f -> (fun () ->
      match l () with
      | x::x2::xs -> twoplus (fun () -> x) (fun () -> x2) (fun () -> xs) s k ()
      | t         -> f (fun () -> t) s k ())
    | Different (empty, one) -> (fun () ->
      match l () with
      | []        -> empty s k ()
      | [x]       -> one (fun () -> x) s k ()
      | x::x2::xs -> twoplus (fun () -> x) (fun () -> x2) (fun () -> xs) s k ())
end

(* Some basic code generators *)
let cunit = fun () -> ()
let update a f = let b = f (liftGet a) in fun () -> a () := b ()
let assign a b = fun () -> (a ()) := (b ())
let apply  f x = fun () -> (f ()) (x ())
let updateM a f = ret (update a f)
let assignM a b = ret (assign a b)
let applyM  f x = ret (apply f x)

let liftRowSwap a b = fun () -> RowSwap (a (), b ())
let liftColSwap a b = fun () -> ColSwap (a (), b ())

(* code transformers *)
module Transformers = struct
    let rec full_unroll lb ub body = 
        if lb>ub then cunit
        else if lb = ub then body lb
        else fun () ->
            begin body lb (); full_unroll (lb+1) ub body () end
end

module Array1Dim = struct
  let init n v = fun() -> Array.init (n ()) (fun _ -> v ())
  let get a i = fun () -> Array.get (a ()) (i ())
  let set a i v = fun () -> Array.set (a ()) (i ()) (v ())
  let length a = fun () -> Array.length (a ())
  let swap a v = fun () -> let (x,y) = v () in
                    let b = a () and t=(a ()).(x) in
                    begin
                        b.(x) <- b.(y);
                        b.(y) <- t;
                        b
                    end
  let blit a1 s1 a2 s2 ln = fun () -> Array.blit (a1 ()) (s1 ()) (a2 ()) (s2 ()) (ln ())
end

(*** tags ****)
module BoolTag = struct
  type 'a tag = 'a * bool

  let tag v b () = v () , b ()
  let get_tag = Pair.snd
  let process v opt f1 f2 = 
    match opt with
    | Some f -> seqM (f (get_tag v))
                     (ifM (fun () -> snd (v ()))
                          (f1 (fun () -> fst (v ())))
                          (f2 ()))
    | None   -> ifM (fun () -> snd (v ()))
                    (f1 (fun () -> fst (v ())))
                    (f2 ())
end

module OptTag = struct
  type 'a tag = 'a option

  let tag v b   = ifL b (Maybe.just (fun () -> v () )) Maybe.none
  let get_tag v = matchOpt v Logic.falseL (fun _ -> Logic.trueL)
  let process v opt f1 f2 = 
    match opt with
      | Some f -> seqM (f (get_tag v)) (matchM v f1 (f2 ()))
      | None   -> matchM v f1 (f2 ())
end

module Vector = struct
  let vec a b () = { vec = a () ; siz = b ()}
  let datum r () = (r ()).vec
  let size r () = (r ()).siz
end