open import GeneralLemmas using (Dec; _≡_)
module DKrivine
  (Node : Set)
  (_≟_ : (n n' : Node)  Dec (n  n'))
  where
open import GeneralLemmas

open import Heap -- new for DKrivine
open import Lambda Node

open import Relation
open import Tagged

RPtr = Ptr × Node -- remote pointer

Closure Env : Set

data StackElem : Set where
  arg : Closure  StackElem
  if0 : Closure  Closure  StackElem
  op₂ : (    )  Closure  StackElem
  op₁ : (  )  StackElem

Closure = Term × Env
ContPtr = RPtr -- new

data EnvElem : Set where
  local  : Closure  EnvElem
  remote : ContPtr    EnvElem -- new; ptr is to a stack in a remote heap; n is the depth of the stack where the closure is

local-inj :  {c₁ c₂}  local c₁  local c₂  c₁  c₂
local-inj refl = refl

remote-inj :  {ptr₁ n₁ ptr₂ n₂}  remote ptr₁ n₁  remote ptr₂ n₂  ptr₁  ptr₂ × n₁  n₂
remote-inj refl = refl , refl

Env = List EnvElem

Stack = List StackElem × Maybe (ContPtr ×  × )
-- at the 'bottom' of the stack we have a remote pointer to a continuation stack along with how many arguments are at the top of that stack (from the offset) and a stack offset
ContHeap = Heap Stack

Thread = Term × Env × Stack -- the analogue of the configuration in the Krivine machine, but possibly not running now
Machine = Maybe Thread × ContHeap -- 'maybe' is running or not; the continuation heap is new

data Msg : Set where -- the messages exchanged on the net: NOTE that the messages are small (important)
  REMOTE : Term     Node  ContPtr    Msg
  -- remote evaluation: the term, the node, the pointer to the sender's cont stack, how many args on that stack
  -- you don't actually need to send the code in an actual implementation because it is known at compile time
  -- sending the code remains an interesting option
  VAR    : ContPtr    ContPtr    Msg
  -- handle variables that are located remotely
  -- remote stack, remote offset, local stack, nr args on local stack
  RETURN : ContPtr      Msg
  -- remote stack, the literal to return, how much to drop from the top of the remote stack

_⊢_▸_ : {A : Set}  Node   Heap A  A  Heap A × RPtr
-- allocation of remote pointer at node i
-- the node, the heap, the value to store in the alloc'ed ptr -> updated heap, remote ptr
i  h  x = (let (h' , ptr) = h  x in h' , (ptr , i))

num-args : Stack  
-- how many arguments are at the top of the stack, including the remote stack if it exists?
-- the remote stack is really a remote extension of the current stack
-- NOTE: this is all LOCAL information about local and remote stacks so to compute it no messages need to be exchanged
num-args ([] , nothing) = 0
num-args ([] , just (contptr , numargs , argstaken)) = numargs
num-args (arg x  s , r) = 1 + num-args (s , r)
num-args (if0 x x₁  s , r) = 0
num-args (op₂ x x₁  s , r) = 0
num-args (op₁ x  s , r) = 0

stack-index : Stack    Maybe EnvElem
-- to evaluate a remote variable, the remote node will use this function when processing a VAR message
-- local stack, an offset -> the local closure or something remote or nothing (that would be an error)
-- NOTE: this is all LOCAL information about local and remote stacks so to compute it no messages need to be exchanged
stack-index (if0 x x₁  s , r) n = nothing
stack-index (op₂ x x₁  s , r) n = nothing
stack-index (op₁ x  s , r) n = nothing
stack-index (arg c  s , r) zero = just (local c)
stack-index (arg c  s , r) (suc i) = stack-index (s , r) i
stack-index ([] , nothing) i = nothing
stack-index ([] , just (contptr , zero , argstaken)) i = nothing
stack-index ([] , just (contptr , suc numargs , argstaken)) zero = just (remote contptr argstaken)
stack-index ([] , just (contptr , suc numargs , argstaken)) (suc i) = stack-index ([] , just (contptr , numargs , suc argstaken)) i

drop-stack : Stack    Maybe Stack
-- drop from the top of a local stack n arguments (error if not enough args)
drop-stack (s , r) zero = just (s , r)
drop-stack ([] , just (contptr , zero , drop)) (suc i) = nothing
drop-stack ([] , just (contptr , suc numargs , drop)) (suc i) = drop-stack ([] , just (contptr , numargs , suc drop)) i
drop-stack ([] , nothing) (suc i) = nothing
drop-stack (arg x  s , r) (suc i) = drop-stack (s , r) i
drop-stack (x  s , r) (suc i) = nothing

-- transitions of DK machine
data _⊢_⟶DKrivine<_>_ (i : Node) : Machine  Tagged Msg  Machine  Set where
  -- Original rules
  -- it knows its node i
  -- just - means the thread is running
  -- r is the remote stack extension
  -- conth is the (new for DK) heap containing continuation stacks (NOTE that heap is quite unstructured, some more structured data type could be used here, to investigate in future work)
  -- the rules are otherwise the same as before, r, conth and i do not change for local executions
  POPARG :  {t e c s r conth} 
    i  (just (ƛ t , e , arg c  s , r) , conth) ⟶DKrivine< silent > (just (t , local c  e , s , r) , conth)
  PUSHARG :  {t t' e s r conth} 
    i  (just ((t $ t') , e , s , r) , conth) ⟶DKrivine< silent > (just (t , e , arg (t' , e)  s , r) , conth)
  VAR :  {n e s r conth t e'}  lookup n e  just (local (t , e')) 
    i  (just (var n , e , s , r) , conth) ⟶DKrivine< silent > (just (t , e' , s , r) , conth)
  COND :  {b t f e s r conth} 
    i  (just (if0 b then t else f , e , s , r), conth) ⟶DKrivine< silent > (just (b , e , if0 (t , e) (f , e)  s , r) , conth)
  COND-0 :  {e t e' f s r conth} 
    i  (just (lit 0 , e , if0 (t , e') f  s , r) , conth) ⟶DKrivine< silent > (just (t , e' , s , r) , conth)
  COND-suc :  {n e t e' f s r conth} 
    i  (just (lit (suc n) , e , if0 t (f , e')  s , r) , conth) ⟶DKrivine< silent > (just (f , e' , s , r) , conth)
  OP :  {f t t' e s r conth} 
    i  (just (op f t t' , e , s , r) , conth) ⟶DKrivine< silent > (just (t , e , op₂ f (t' , e)  s , r) , conth)
  OP₂ :  {n e f t e' s r conth} 
    i  (just (lit n , e , op₂ f (t , e')  s , r) , conth) ⟶DKrivine< silent > (just (t , e' , op₁ (f n)  s , r) , conth)
  OP₁ :  {n e f s r conth} 
    i  (just (lit n , e , op₁ f  s , r) , conth) ⟶DKrivine< silent > (just (lit (f n) , [] , s , r) , conth)

  -- New rules
  -- OBS: This is 'purely functional' in the sense that the remote stack extensions don't need to change, just local views
     -- This means that should make it easier to make parallelism because you can have several views of the same stack without conflicts
  POPARG-remote :  {t e contptr args drop conth} 
    i  (just (ƛ t , e , [] ,                     just (contptr , suc args , drop)) , conth) ⟶DKrivine< silent >
        (just (t , remote contptr drop  e , [] , just (contptr , args , suc drop)) , conth)
    -- local stack is empty so we need to use the remote stack extension. it needs to have at least one argument.
    -- the environment is updated with the 'view' of the remote pointer
    -- the view of the remote stack extension is also updated
    -- OBS: This is not sending anything, just changes the local view of the remote stack

  REMOTE-send :  {t i' e s conth}  let (conth' , contptr) = i  conth  s in
    i  (just (t at i' , e , s) , conth) ⟶DKrivine< send (REMOTE t i' contptr (num-args s)) >
        (nothing , conth')
    -- remote execution:
    -- we create a pointer to the current stack in the heap (contptr)
    -- send the remote exec message
    -- machine halts

  REMOTE-receive :  {conth t contptr numargs} 
    i  (nothing , conth) ⟶DKrivine< receive (REMOTE t i contptr numargs) > (just (t , [] , [] , just (contptr , numargs , 0)) , conth)
    -- the machine was halting
    -- start exec with current heap of cont stacks
    -- env is empty (only closed terms)
    -- at the bottom we have the remote stack extension from the origin node of the incoming msg

  -- this is the trickiest set of rules!
  VAR-send :  {n e s rcontptr index conth}  lookup n e  just (remote rcontptr index)  let (conth' , contptr) = i  conth  s in
    i  (just (var n , e , s) , conth) ⟶DKrivine< send (VAR rcontptr index contptr (num-args s)) > (nothing , conth')
    -- we reach a remote var (we can see in the env it is remote)
    -- we save our current cont stack in the heap, pointed at by contptr
    -- send a var message (rcontptr includes the node)
    -- halt

  VAR-receive :  {conth contptr s n rcontptr numargs el}  conth ! contptr  just s  stack-index s n  just el 
    i  (nothing , conth) ⟶DKrivine< receive (VAR (contptr , i) n rcontptr numargs) >
        (just (var 0 , el  [] , [] , just (rcontptr , numargs , 0)) , conth)
    -- the halting machine gets the remote var message
    -- contptr points into the local heap (at i)
    -- rcontptr points into the origin node of the message
    -- n is the index into the stack at i (of contptr)
    -- NOTE: we avoid having special cases on whether index n in contptr is local or remote.
       -- just el can be local or remote
       -- we create var 0 as a 'dummy' variable which refers to el in the environment this may trigger the VAR-send again
       -- or the (local) VAR is the case may be.

  RETURN-send :  {n e contptr drop conth} 
    i  (just (lit n , e , [] , just (contptr , 0 , drop)) , conth) ⟶DKrivine< send (RETURN contptr n drop) > (nothing , conth)
    -- we reached a literal, we send it back
  RETURN-receive :  {conth contptr s s' n drop}  conth ! contptr  just s  drop-stack s drop  just s' 
    i  (nothing , conth) ⟶DKrivine< receive (RETURN (contptr , i) n drop) > (just (lit n , [] , s') , conth)
    -- receive a literal, the cont stack we use is the one from contptr after we dropped drop elements

open import Network Node _≟_ _⊢_⟶DKrivine<_>_ public

initial-network-sync : Term  Node  SyncNetwork
initial-network-sync t i = update  i'  (nothing , )) i (just (t , [] , [] , nothing), )

initial-network-async : Term  Node  AsyncNetwork
initial-network-async c i = initial-network-sync c i , []