open StateCPSMonad

module type Value = sig
  type value
end

module IntV = struct
  type value = int
end

module Traverse (CODE: Coderep.T) = struct

  open CODE

  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

  module ContainerLib (V:Value) (L:Value) = struct
    module type Sig = sig
      type value = V.value
      type 'a container
      type eltdesc
      type loopdata = L.value

      val tr : value container -> value container

      val access : eltdesc abstract -> value abstract
      val traverseexercise :
        value container abstract ->
        (eltdesc abstract -> loopdata abstract -> (value * loopdata) abstract) ->
        loopdata abstract ->
        ('a, value container) cmonad
    end
  end

  module ListContainerLib (V : Value) (L : Value) = struct
    type value = V.value
    type 'a container = 'a list
    type eltdesc = value
    type loopdata = L.value

    let tr l =
    let rec tr_ l i =
    match l with
    |[] -> []
    |_::t -> i::(tr_ t (i+1))
    in tr_ l 0

    let access x = x

    let traverseexercise lst f data  =
      let gen self tup =
      let (lst, i) = Pair.unpair tup in
      CList.matchListM lst
        (ret CList.nil)
        (fun h t ->
        let (v, i2) = Pair.unpair (f (h) i) in
        let! r = (applyM self (Tuple.tup2 t i2)) in
        ret (CList.cons v r))
      in genrecloop gen (Tuple.tup2 lst data)
  end

  module ArrayContainerLib (V : Value) (L : Value) = struct
    type value = V.value
    type 'a container = 'a array
    type eltdesc = value container * int
    type loopdata = L.value

    let tr a =
      let i = ref 0 in
      let len = Array.length a in
      while !i < len do
        Array.set a !i !i;
        i := !i + 1
      done;
      a

    let access x =
      let (arr, n) = Pair.unpair x in
      Array1Dim.get arr n

    let traverseexercise arr f data =
      let! ix = retN (liftRef Idx.zero) in
      let i = liftGet ix in
      let! nr = retN (liftRef data) in
      let n = liftGet nr in
      let cond = Idx.less i (Array1Dim.length arr) in
      seqM
        (whileM cond (
          let (v, n2) = Pair.unpair (f (Tuple.tup2 arr i) n) in
          seqM (seqM
            (ret (Array1Dim.set arr i v))
            (assignM nr n2))
            (assignM ix (Idx.succ i))))
        (retN arr)
  end

  module Ex (Container : ContainerLib (IntV) (IntV).Sig) = struct
    let traverseexercise container =
      let body = fun _ n -> Tuple.tup2 n (Idx.succ n) in
      Container.traverseexercise container body Idx.zero
  end

end

module InstantiateCode = struct
  open Runcode
  open Print_code

  let instantiate gen =
  let code = .< fun a -> .~(runM (gen .<a>.) []) >. in
  let ppf = Format.std_formatter in
      print_endline "before-test";
      let (cde, _) = close_code_delay_check code in
      Format.fprintf ppf ".<@,%a>.@ " format_code cde;
      print_newline ();
      print_endline "after-test";
      Runcode.run code
end

module Test = struct
  module TC = Traverse(Code)
  module LCII = TC.ListContainerLib (IntV) (IntV)
  module ACII = TC.ArrayContainerLib (IntV) (IntV)
  module LT = TC.Ex (LCII)
  module AT = TC.Ex (ACII)
  let test _ =
    let listtest a = LT.traverseexercise a in
    let arraytest a = AT.traverseexercise a in
    let resultl = InstantiateCode.instantiate listtest [9;8;7;5] in
    let resulta = InstantiateCode.instantiate arraytest (Array.init 20 (fun x -> x*2)) in
    assert (resultl = [0;1;2;3]);
    assert (resulta = (Array.init 20 (fun x -> x)))
end;;

Test.test ();;