/* File:        bangs.er
   Author:      John Longley
   Date:        March 2009
   Description: Universal languages for the Lamarche model and categories 
                arising from various bangs on it.
*/

/* The universal type: */

type forest = rectype f => int -o {label: int, rest: f} ;
fun hit ((fold forest F) :forest) = oncefn i:int => F $ i ;


/* Coding int+int as int */

fun inL (i:int) = i+i ;
fun inR (i:int) = (i+i)+1 ;
recfun halve (i:int) :int = if i<2 then 0 else halve (i-2) + 1 ;
fun outL (i:int) = if ~1<i then halve i else 0-halve(1-i) ;
val outR = outL ;
recfun even (i:int) :bool = if i<2 then i=0 else even (i-2) ;
fun isL (i:int) = if ~1<i then even i else even (0-i) ;

/* Coding operations for integer lists (used as node addresses in trees) */

type intlist = rectype t => [| nil, cons of {hd:int, tl:t} |] ;
val nil = fold intlist (tok nil) ;
fun cons (i:int) (il:intlist) = fold intlist (tag cons {hd=i, tl=il}) ;
fun isPrefix (il:intlist) (il':intlist) = true ; // dummy
fun labelWithNode (i:int) (il:intlist) = i ;     // dummy 
fun labelOf (i:int) = i ;                        // dummy
fun dataOf (i:int) = i ;                         // dummy
fun packList (il:intlist) = 0 ;                  // dummy
fun unpackList (i:int) = nil ;                   // dummy 


/* State operations */

/* Integer store cells (borrowed from class-sugar.er): */

type intCell = {get: unit->int, set: int->unit} ;

val intCellClass =
   class :: intCell with
      fields i:int
      methods
         get () = ?i ,
         set j  = i:=j
      constr j:int => {i=j}
   end ;

val mkIntCell = new intCellClass ;

mkIntCell :< int -> intCell ;


/* Trace operation for hiding internal state of type "forest" */

type intForest  = {i:int, f:forest} ;
type traceArgs = {initial: forest, op: !(intForest -o intForest)} ;

type forest' = {f:forest} ;
type intForestVS  = {value:int, state:forest'} ;
type intForestVS' = {value:!int, state:forest'} ;
type forestSelfTy = {m: intForestVS -> intForestVS'} ;

val trace =
   fn {initial=initial, op=op} : traceArgs =>
   make (linear extend linear root with   
        fn super: {} => fn self: forestSelfTy =>
        // will look nicer when sugared!
           {m = fn (i \ F) : intForestVS =>
                split op {i=i, f=F.f} as {i=j,f=g}
                in prom j \ {f=g}
                end} ,
        fn supcon: {}->{} => fn f0:forest => {f=f0}
        end) 
        initial 
        .m ;

trace :< traceArgs -> (int -> int) ;

/* More powerful version with result type "tree".
   Used below in pseudopromotion. 
   Lives in all the models that have "trace". */

type tree = {label:int, rest:forest} ;

type intForest  = {i:int, f:forest} ;
type treeForest = {t:!tree, f:forest} ;  
// the bang here is needed in the backtracking setting, but not in
// connection with the simple "repetitive exponential".
type treeTraceArgs = {initial: forest, op: !(intForest -o treeForest)} ;
type treeForestVS = {value:!tree, state:forest'} ;
type treeForestSelfTy = {m: intForestVS -> treeForestVS} ;

val treeTrace =
   fn {initial=initial, op=op} : treeTraceArgs =>
   make (linear extend linear root with   
        fn super: {} => fn self: treeForestSelfTy =>
           {m = fn (i \ F) : intForestVS =>
                split op {i=i, f=F.f} as {t=t,f=g}
                in t \ {f=g}
                end} ,
        fn supcon: {}->{} => fn f0:forest => {f=f0}
        end) 
        initial 
        .m ;

treeTrace :< treeTraceArgs -> (int -> tree) ;


/* Exactly the same thing, but with a more complicated type subtrees'
   (this illustrates why we want polymorphism in the language!).
   Used below to implement "mergeSubtrees". */

type subtreeData = {label:int, daughters:!forest} ;
type subtrees = [| firstTime of int -o subtreeData, 
                   notFirstTime |] ;

type subtrees' = {s:subtrees} ;
type intSubtrees  = {i:int, s:subtrees} ;
type treeSubtrees = {t:!tree, s:subtrees} ;  
type treeTrace2Args = 
      {initial: subtrees, op: !(intSubtrees -o treeSubtrees)} ;
type intSubtreesVS  = {value:int, state:subtrees'} ;
type treeSubtreesVS = {value:!tree, state:subtrees'} ;
type treeSubtreesSelfTy = {m: intSubtreesVS -> treeSubtreesVS} ;

val treeTrace2 =
   fn {initial=initial, op=op} : treeTrace2Args =>
   make (linear extend linear root with   
        fn super: {} => fn self: treeSubtreesSelfTy =>
           {m = fn (i \ S) : intSubtreesVS =>
                split op {i=i, s=S.s} as {t=t,s=g}
                in t \ {s=g}
                end} ,
        fn supcon: {}->{} => fn s0:subtrees => {s=s0}
        end) 
        initial 
        .m ;

treeTrace2 :< treeTrace2Args -> (int -> tree) ;


/* Weaker stateful operation: uses a single forest as a script for 
   two forests side by side. Implemented here using trace.

   A more "primitive" programming construct may be possible here, 
   but probably more complication than it is worth.
*/

type intfn = int -> int ;

fun forestToIntfn (f:forest) =
    trace {initial = f,
           op = fn {i=i, f=fold forest F} : intForest =>
                (val {label=j, rest=g} = F$i ; {i=j, f=g})} ;

fun IntfnToForestL (h:intfn) =
    rec f:forest =>
        fold forest (oncefn i:int => 
	    {label = h (inL i), rest = f}) ;

fun IntfnToForestR (h:intfn) =
    rec f:forest =>
        fold forest (oncefn i:int => 
	    {label = h (inR i), rest = f}) ;

fun shareForest (f:forest) =
   (val h = forestToIntfn f ; 
    {L = IntfnToForestL h, R = IntfnToForestR h}) ;

shareForest :< forest -> {L: forest, R: forest} ;

 
/* Embedding forest in !forest (diverging after first question 
   at each level). Implemented using treeTrace. */

val bottomForest = rec f:!forest => f ;

recfun pseudopromote (f0:forest) :!forest = 
    (val branching =
      treeTrace {initial=f0,
        op = fn {i=i, f=fold forest F} : intForest =>
             split F $ i as {label=j, rest=g} in 
                (val g' = pseudopromote g ;
                 {t = prom {label=j, 
                            rest=fold forest (oncefn k:int => 
                                 unfold (der g') $ k)},
                  f = der bottomForest})
             end} ;
     prom (fold forest (oncefn k:int => branching k))) ;

pseudopromote :< forest -> !forest ;


/* Using this, we can implement the "once only" version of catchcont
   (much as in theory.er): */

type onceCatchcontType =
   [| result of {value  : int , 
                 more   : (int -o tree) -o forest} ,
      query  of {arg    : int ,
                 resume : tree -o {value: int, residue: forest}} |] ;

val bottomBranching = rec T: int->tree => T ;

fun pseudopromoteBr (br:int -o tree) =
    (val F = pseudopromote (fold forest br) ;
     fn i:int => unfold (der F) $ i) ;

fun pseudopromoteTree ({label=i, rest=f} : tree) =
    (val f' = pseudopromote f ;
     prom {label = i, 
           rest = fold forest (oncefn j:int => unfold (der f') $ j)}) ;

type intForestVR = {value:int, residue:forest} ;

fun onceCatchcont (body : (int -o tree) -o intForestVR) =
        case catchcont x: !(int -o tree) => body $ (der x) of
            result r =>
                 split r as {value=v, more=M} in
                    tag result 
                       {value = v, 
                        more = oncefn x : int -o tree => 
                                 M $ (pseudopromoteBr x)}
                    :< onceCatchcontType
                 end
          | query q =>
                 split q as {arg=a, resume=R} in
                    tag query
                       {arg = a,
                        resume = oncefn t:tree => 
                           R $ (pseudopromoteTree t) $ bottomBranching}
                    :< onceCatchcontType
                 end
        end ;

// Note: it's a runtime fact about onceCatchcont that bottomBranching 
// and bottomForest are never invoked.

onceCatchcont :< ((int -o tree) -o intForestVR) -> onceCatchcontType ;


/* Implementation of critical retractions for the basic (bang-free) model.
   We require just onceCatchcont and shareForest. */

/* For product */

val unpair = shareForest ;

/* For "-o" */

type tree = {label:int, rest:forest} ;
type branching = int -o tree ;

type abstractBodyType = {
    fst : (branching -o forest) -> forest ,
    snd : (branching -o tree) -> tree
    } ;

val FF = fn b : branching => fold forest b ;

fun abstract (f: forest -o forest) =
   (rec abstr : abstractBodyType =>
          {fst = fn P : branching -o forest =>
	         FF (oncefn h:int => 
	             abstr.snd (oncefn g : branching => 
	                        (unfold (P$g)) $ (outR h))) ,
	   snd = fn R : branching -o tree =>
		 case 
                      onceCatchcont (oncefn x: int -o tree => 
			 split R $ x as {label=l, rest=r} in
                            {value=l, residue=r} 
                         end)
		 of
		      result r => 
                         split r as {value=v, more=M} in 
                             {label = inR v , rest = abstr.fst M} 
                         end
		    | query q => 
                         split q as {arg=a, resume=R} in 
			     {label = inL a ,
			      rest = FF (oncefn h:int =>  
					 abstr.snd (oncefn b:branching =>
                                             split R $ {label= outL h,
                                                        rest = (FF b)}
                                             as {value=l, residue=r} 
                                             in {label=l, rest=r}
                                             end))}
	                 end 
		 end})
   .fst (oncefn b: branching => f $ (FF b)) ;

type applyBodyType = {
    fst : forest -> (forest -o forest) ,
    snd : int -> forest -o (forest -o tree)
    } ;

val apply = (rec app:applyBodyType =>
   {fst = fn f:forest => oncefn g:forest =>
        fold forest (oncefn i:int => 
	    split hit f $ (inL i) as {label=j, rest=f'} in
                if isL j then {label = outL j, rest = app.fst f' $ g}
                else app.snd (outR j) $ f' $ g
            end) ,
    snd = fn k:int => oncefn f:forest => oncefn g:forest =>
        split hit g $ k as {label=l, rest=g'} in
            split hit f $ (inR l) as {label=j, rest=f'} in
                if isL j then {label = outL j, rest = app.fst f' $ g'}
                else app.snd (outR j) $ f' $ g'
            end
        end})
   .fst ;

abstract :< (forest -o forest) -> forest ;
apply :< forest -> (forest -o forest) ;


/* Next we exhibit !forest as a retract of forest, for each of the
   bangs of interest. */

/* An operation using "force" used in the embedding !forest -> forest
   for bangs 1, 3 and 4. */

fun forceApply (f:!forest) (m:int) = 
       (val p = force (fn u:unit => 
                   split (unfold (der f) $ m) as {label=n, rest=f'} 
                   in {value=n, residue=f'} end) ;
        {value = der p .value, 
	 residue = prom (fold forest (oncefn i:int => 
                   unfold (der p.residue) $ i))}) ;

forceApply :< !forest -> int -> {value: int, residue: !forest} ;

// Syntactic note: perhaps "fprom e" is nicer than "force (fn u => e)"


/* ! 1: Lamarche exponential (non-repetitive backtracking). */

/* Embedding: requires only forceApply */

recfun extract1 (subforests: int->!forest) (currNode: intlist) : forest =
    fold forest (oncefn m:int =>
       (val {value=n, residue=G} = 
            forceApply (subforests (labelOf m)) (dataOf m) ;
        val newNode = cons m currNode ;
        val newLabel = packList newNode ;
        {label = n,
         rest = extract1 
                  (fn i:int => if i=newLabel then G else subforests i)
                  newNode})) ;

val nilLabel = packList nil ;

fun embed1 (F: !forest) = 
    extract1 (fn i:int => if i=nilLabel then F else rec G:!forest => G) 
             nil ;

/* Projection: requires shareForest plus a new operation mergeSubtrees */

// Recall:
type subtreeData = {label:int, daughters:!forest} ;
type subtrees = [| firstTime of int -o subtreeData, 
                   notFirstTime |] ;
type intSubtrees  = {i:int, s:subtrees} ;
type treeSubtrees = {t:!tree, s:subtrees} ;  
type treeTrace2Args = 
      {initial: subtrees, op: !(intSubtrees -o treeSubtrees)} ;

// To implement mergeSubtrees, we use treeTrace2 and forceApply

type subtreesData = {first : int -o subtreeData, others : !forest} ;

fun mergeSubtrees ({first=first, others=others} : subtreesData) =
    (val F = treeTrace2
        {initial = (tag firstTime first) :< subtrees,
         op = fn {i=i, s=s} :intSubtrees => case s of
                  firstTime st =>
                     (val {label=j, daughters=fd} = st $ i ;
                      {t = prom {label=j, rest=fold forest (oncefn k:int => hit (der fd) $ k)} , 
                       s = (tok notFirstTime) :< subtrees}) 
                | notFirstTime =>
                     (val {value=j, residue=fs} = forceApply others i ;
                       // can we dispense with forceApply here?
                      {t = prom {label=j, rest=fold forest (oncefn k:int => hit (der fs) $ k)} , 
                       s = (tok notFirstTime) :< subtrees}) 
                  end } ;
    prom (fold forest (oncefn k:int => F k))) ;

mergeSubtrees :< subtreesData -> !forest ;

/* Added later: I think mergeSubtrees could be implemented using 
   a simpler primitive
      switch_t : (int -o (int -> t)) -> (int -> t)
   for which e.g. an int store cell suffices.
*/


// Projection, using shareForest and mergeTrees:

recfun treeDiverge (_:unit) : tree = treeDiverge () ;

recfun partition' (node:intlist) (f:forest) : forest = 
    fold forest (oncefn i:int =>
        if isL i then
           if isPrefix node (unpackList (labelOf (outL i)))
           then (val {label=j, rest=f'} = hit f $ (outL i) ;
                 {label=j, rest = partition' node f'})
           else treeDiverge()
        else
           if isPrefix node (unpackList (labelOf (outR i)))
           then treeDiverge()
           else (val {label=j, rest=f'} = hit f $ (outR i) ;
                 {label=j, rest = partition' node f'})) ;

fun partition (node:intlist) (f:forest) =
    fold forest (oncefn i:int =>
        if isL i then
           if unpackList (labelOf (outL i)) = node 
           then (val {label=j, rest=f'} = hit f $ (outL i) ;
                 {label=j, rest = partition' (cons i node) f'})
           else treeDiverge()
        else treeDiverge()) ;

recfun project1aux (node:intlist) (f:forest) : !forest = 
   (val {L=f1, R=f2} = shareForest (partition node f) ;
    mergeSubtrees 
       {first = oncefn i:int =>
           (val {label=j, rest=f1'} = hit f1 $ (labelWithNode i node) ;
            {label = j, daughters = project1aux (cons i node) f1'}) ,
  	others = project1aux node f2}) ;

val project1 = project1aux nil ;

embed1 :< !forest -> forest ;
project1 :< forest -> !forest ;

/* Co-Kleisli category for ! 1 */

/* In comments we'll write --> for the Kleisli arrow: A --> B is (!A) -o B.
   By contrast, note that in Lingay, A -> B is !(A -o B) .)

   We can still use "forest" as our universal type, since !int ~ int
   so int --> A ~ int -o A.

   For ease of doing pure "FPC programming" relative to our primitives, 
   we'll implement all our primitives in promoted form,
   i.e. with types !A1 -o ... -o !Am -o !B1 * ... * !Bn (curried)
                or !A1 *  ...  * !Am -o !B1 * ... * !Bn (uncurried)
   building in a use of "force" where necessary.
*/

fun hit' (F : !forest) (i : int) = hit (der F) $ i ;

fun pair1 (F1 : !forest) (F2 : !forest) =
    prom (fold forest (oncefn n:int =>
         if isL n then hit' F1 (outL n) else hit' F2 (outR n))) ;

fun unpair1 (F : !forest) =
   {fst = prom (fold forest (oncefn n:int => hit' F (inL n))),
    snd = prom (fold forest (oncefn n:int => hit' F (inR n)))} ;

/* For -->, it seems what we need is exactly "catchcopy".
   Strangely, the beast we want here is none other than the Lingay "catchcont".
   But it has a different "meaning" in the coKleisli world, where composition
   requires "the use of force" ;>)
   Here, for convenience, we build in the distribution of ! over +.
*/

type resultType = {value  : int , 
                   more   : (int -> tree) -o forest} ;

type queryType =  {arg    : int ,
                   resume : !tree -o (int -> tree) -o intForestVR} ;

/* Recall: type intForestVR = {value: int, residue: forest} ; */

type catchcontType = [| result of resultType , query of queryType |] ;

type catchcopyType =
   [| result of {value  : int , 
                 more   : (int -> tree) -> forest} ,
      query  of {arg    : int ,
                 resume : !tree -> (int -> tree) -> intForestVR} |] ; 

recfun dummyResult (_:unit) : resultType = dummyResult () ;

recfun dummyQuery (_:unit) : queryType = dummyQuery () ;

type bangCatchcontType = [| result of !resultType, query of !queryType |] ;

fun distrib (X : !catchcontType) =
    case der X of
       result r' => tag result (force (fn u:unit => case der X of 
	                result r => r | query q => dummyResult() end ))
                    :< bangCatchcontType
     | query q'  => tag query (force (fn u:unit => case der X of 
                        query q => q  | result r => dummyQuery() end ))
                    :< bangCatchcontType 
    end ;

fun catchcopy (body: (int -> tree) -> intForestVR) =
    case distrib (force (fn u:unit => 
	 catchcont test : int->tree => body test)) of
       result R => tag result 
                      {value = der R .value ,
                       more = fn dummy : int->tree => der R .more $ dummy} 
                   :< catchcopyType
     | query Q  => tag query
                      {arg = der Q .arg ,
                       resume = fn t : !tree => fn dummy : int->tree =>
                             der Q .resume $ t $ dummy}
                   :< catchcopyType 
    end ;

catchcopy :< ((int -> tree) -> intForestVR) -> catchcopyType ;


/* Aside: It would be nice to code up "catchcont" using the complete 
   set of operators for the linear category given above.
   Unfortunately, it seems that catchcont isn't definable from
   onceCatchcont uniformly in the type, so we'd need to go via
   some unpleasant type-specific retractions...  
*/



/* The acid test is now whether we can code 
       abstract1 : (forest --> forest) --> forest
   with just catchcopy. (Essentially the "Stratagem program"!)
   It might take some effort to impose on ourselves the discipline of
   programming entirely "within the coKleisli model".
*/

fun VRtoLR ({value=l, residue=r} : intForestVR) = {label=l, rest=r} :< tree ;
fun LRtoVR ({label=l, rest=r} : tree) = {value=l, residue=r} :< intForestVR ;

type abstract1BodyType = {
    fst : (!forest -> forest) -> forest ,
    snd : intlist -> ((int->tree) -> (int->tree) -> tree) -> tree ,
    thd : ((int->tree) -> forest) -> forest ,
    fth : catchcopyType -> tree
    } ;

val abstract1 =
   (rec abstr : abstract1BodyType =>
        {fst = fn P : !forest -> forest => 
               FF (oncefn h:int =>
                   abstr.snd nil
                       (fn x: int->tree => fn g: int->tree =>
                            (unfold (P (prom (fold forest (oncefn i:int => g i))))) 
                             $ (outR h))) ,
         snd = fn node: intlist =>
               fn R : (int->tree) -> (int->tree) -> tree => 
                   abstr.fth (catchcopy (fn x: int->tree =>
                       LRtoVR (R x (fn i:int => x (labelWithNode i node))))) ,
         thd = fn Q : (int->tree) -> forest => 
               FF (oncefn h:int =>
                   abstr.fth (catchcopy (fn x: int->tree =>
                       LRtoVR (unfold (Q x) $ (outR h))))) ,
         fth = fn stuff: catchcopyType =>
               case stuff of
                   result {value=v, more=M} =>
                      {label = inR v, rest = abstr.thd M}
                 | query {arg=a, resume=Z} =>
                      {label = inL a, 
                       rest = FF (oncefn h:int =>
                              (val h' = outL h ;
                               abstr.snd (cons h (unpackList (labelOf a)))
                                   (fn x: int->tree => fn b: int->tree =>
                                       VRtoLR (Z (prom {label = h', 
                                                        rest = fold forest 
                                                               (oncefn i:int => b i)})
                                                 x))))}
               end })
   .fst ;

abstract1 :< !(!forest -o forest) -> forest ;
// Phew!




/* Old stuff concerning another universal type "bangforest".
   A red herring I think, arising from confusion between --> and -> */

/*
type bangforest = rectype f => int -> {label: int, rest: f} ;

// !forest as a retract of bangforest

recfun deepDerelict ((fold bangforest B) : bangforest) : forest = 
   fold forest (oncefn i:int => 
      (val {label=j, rest=B'} = B i ;
       {label=j, rest=deepDerelict B'})) ;

fun innerDerelict (B : bangforest) =
   prom (fold forest (oncefn i:int => hit (deepDerelict B) $ i)) ;

recfun deepForce (F: !forest) : bangforest =
   fold bangforest (fn i:int =>
      (val {value=j, residue=F'} = forceApply F i ;
       {label=j, rest=deepForce F'})) ;

innerDerelict :< bangforest -> !forest ;
deepForce :< !forest -> bangforest ;
*/