# Tiny UNIX revision 2.

typename Option(a) = [|None|Some:a|];

sig fail : () {Fail:Zero |_}-> a
fun fail() { switch (do Fail) {} }

sig optionalise : (Comp(a, {Fail:Zero |e})) {Fail{_} |e}~> Option(a)
fun optionalise(m) {
  handle(m()) {
    case Return(x) -> Some(x)
    case Fail(_)   -> None
  }
}

sig withDefault : (a, Comp(a, {Fail:Zero |e})) {Fail{_} |e}~> a
fun withDefault(x', m) {
  handle(m()) {
    case Return(x) -> x
    case Fail(_)   -> x'
  }
}

sig lookup : (a, [(a, b)]) {Fail:Zero |_}~> b
fun lookup(k, kvs) {
  switch (kvs) {
    case [] -> fail()
    case (k', v) :: kvs' ->
      if (k == k') v
      else lookup(k, kvs')
  }
}

sig modify : (a, b, [(a, b)]) ~> [(a, b)]
fun modify(k, v, kvs) {
  switch (kvs) {
    case [] -> []
    case (k', v') :: kvs' ->
      if (k == k') (k, v) :: kvs'
      else (k', v') :: modify(k, v, kvs')
  }
}

sig remove : (a, [(a, b)]) ~> [(a, b)]
fun remove(k, kvs) {
  switch (kvs) {
    case [] -> []
    case (k', v') :: kvs' ->
      if (k == k') kvs'
      else (k', v') :: remove(k, kvs')
  }
}

sig has : (a, [(a, b)]) ~> Bool
fun has(k, kvs) {
  switch (kvs) {
    case [] -> false
    case (k', _) :: kvs' ->
      k == k' || has(k, kvs')
  }
}

##
## Basic i/o
##
typename File = String;
typename FileDescr = Int;

sig stdout : FileDescr
var stdout = 1;

sig basicIO : (Comp(a, {Write:(FileDescr,String) -> () |e})) {Write{_} |e}~> (a, File)
fun basicIO(m) {
  handle(m()) {
    case Return(res) -> (res, "")
    case Write(_, s, resume) ->
      var (res, file) = resume(());
      (res, s ^^ file)
  }
}

sig echo : (String) {Write:(FileDescr,String) -> () |_}~> ()
fun echo(cs) { do Write(stdout,cs) }

fun example0() {
  basicIO(fun() {
        echo("Hello"); echo("World")
  })
}

##
## Exceptions: non-local exits
##
sig exit : (Int) {Exit:(Int) -> Zero |_}-> a
fun exit(n) { switch(do Exit(n)) {} }

sig status : (Comp(a, {Exit:(Int) -> Zero |e})) {Exit{_} |e}~> Int
fun status(m) {
  handle(m()) {
    case Return(_) -> 0
    case Exit(n, _) -> n
  }
}

fun example1() {
  basicIO(fun() {
        status(fun() {
          echo("dead"); exit(1); echo("code")
        })
  })
}

##
## Dynamic binding: user-specific environments.
##
typename User = [|Alice|Bob|Root|];

sig whoami : () {Ask:String |_}-> String
fun whoami() { do Ask }

sig env : (User, Comp(a, {Ask:String |e})) {Ask{_} |e}~> a
fun env(user, m) {
  handle(m()) {
    case Return(x) -> x
    case Ask(resume) ->
      switch (user) {
        case Alice -> resume("alice")
        case Bob   -> resume("bob")
        case Root  -> resume("root")
      }
  }
}

fun example2() {
  env(Root, whoami)
}

###
### Session management.
###
sig su : (User) {Su:(User) -> () |_}-> ()
fun su(user) { do Su(user) }

sig sessionmgr : (User, Comp(a, {Ask:String,Su:(User) -> () |e})) {Ask{_},Su{_} |e}~> a
fun sessionmgr(user, m) {
  env(user, fun() {
     handle(m()) {
       case Return(x) -> x
       case Su(user', resume) ->
         env(user', fun() { resume(()) })
     }
  })
}

fun example3() {
  basicIO(fun() {
        sessionmgr(Root, fun() {
          status(fun() {
            su(Alice); echo(whoami()); echo(" ");
            su(Bob);   echo(whoami()); echo(" ");
            su(Root);  echo(whoami())
          })
        })
  })
}

##
## Nondeterminism: time sharing.
##
sig fork : () {Fork:Bool |_}-> Bool
fun fork() { do Fork }

sig nondet : (Comp(a, {Fork:Bool |e})) {Fork{_} |e}~> [a]
fun nondet(m) {
  handle(m()) {
    case Return(res) -> [res]
    case Fork(resume) -> resume(true) ++ resume(false)
  }
}

sig interrupt : () {Interrupt:() |_}-> ()
fun interrupt() { do Interrupt }

typename Pstate(a,e::Eff) = forall q::Presence.
                            [|Done:a
                             |Paused:() {Interrupt{q} |e}~> Pstate(a, { |e})|];

sig slice : (Comp(a, {Interrupt:() |e})) {Interrupt{_} |e}~> Pstate(a, { |e})
fun slice(m) {
  handle(m()) {
    case Return(res) -> Done(res)
    case Interrupt(resume) -> Paused(fun() { resume(()) })
  }
}

sig schedule : ([Pstate(a, { Fork:Bool |e})]) {Fork{_},Interrupt{_} |e}~> [a]
fun schedule(ps) {
  # sig run : (Pstate (a, { Fork:() {}-> Bool|e })) {Fork{_},Interrupt{_} |e}~> [a]
  # fun run(p) {
  #   switch(p) {
  #     case Done(res) -> [res]
  #     case Paused(resume) -> runNext(nondet(resume))
  #   }
  # }
  # concatMap(run, ps)
  fun schedule(ps, done) {
    switch (ps) {
      case [] -> done
      case Done(res) :: ps' ->
        schedule(ps', res :: done)
      case Paused(resume) :: ps' ->
        schedule(ps' ++ nondet(resume), done)
    }
  }
  schedule(ps, [])
}

sig timeshare : (Comp(a, {Fork:Bool,Interrupt:() |e})) {Fork{_},Interrupt{_} |e}~> [a]
fun timeshare(m) {
  var p = Paused(fun() { slice(m) });
  schedule([p])
}

fun example4() {
  basicIO(fun() {
     timeshare(fun() {
        sessionmgr(Root, fun() {
           status(fun() {
             var parent = fork();
             (if (parent) su(Alice) else su(Bob));
             echo(whoami() ^^ "> Hello ");
             interrupt();
             echo(whoami() ^^ "> Bye ");
             interrupt();
             if (parent) exit(0)
             else {
               var parent = fork();
               interrupt();
               (if (parent) su(Root)
                else {
                 echo(whoami() ^^ "> oops ");
                 exit(1)
               });
               echo(whoami() ^^ "> !! ")
             }
           })
        })
     })
  })
}

fun forktest(n) {
  fun loop(i, n) {
    if (i >= n) (-1)
    else {
      var x = fork();
      println("< x = " ^^ (if (x) "true" else "false") ^^ ", i = " ^^ intToString(i));
      ignore(loop(i+1,n));
      println("> x = " ^^ (if (x) "true" else "false") ^^ ", i = " ^^ intToString(i));
      exit(i)
    }
  }
  loop(0, n)
}

fun test(i) {
  if (i == 2) ()
  else {
    println("< i = " ^^ intToString(i));
    var x = fork();
    test(i+1);
    println("> i = " ^^ intToString(i))
  }
}

fun ritchie() {
  echo("UNIX is basically ");
  echo("a simple operating system, ");
  echo("but ");
  echo("you have to be a genius to understand the simplicity.\n")
}

fun hamlet() {
  echo("To be, or not to be,\n");
  echo("that is the question:\n");
  echo("Whether 'tis nobler in the mind to suffer\n")
}

fun example5() {
  basicIO(fun() {
    nondet(fun() {
      sessionmgr(Root, fun() {
        status(fun() {
          if (fork()) {
            su(Alice);
            ritchie()
          } else {
            su(Bob);
            hamlet()
          }
        })
      })
    })
  })
}

fun interruptWrite(m) {
  handle(m()) {
    case Return(res) -> res
    case Write(fd, cs, resume) ->
      interrupt(); resume(do Write(fd, cs))
  }
}

fun example5'() {
  basicIO(fun() {
    timeshare(fun() {
      interruptWrite(fun() {
        sessionmgr(Root, fun() {
          status(fun() {
            if (fork()) {
              su(Alice);
              ritchie()
            } else {
              su(Bob);
              hamlet()
            }
          })
        })
      })
    })
  })
}

##
## Generic state handling
##
sig get : () {Get:s |_}-> s
fun get() { do Get }

sig put : (s) {Put:(s) -> () |_}-> ()
fun put(st) { do Put(st) }

sig runState : (s, Comp(a, {Get:() -> s,Put:(s) -> () |e})) {Get{_},Put{_} |e}~> (a, s)
fun runState(st0, m) {
  var f = handle(m()) {
    case Return(x) -> fun(st) { (x, st) }
    case Get(resume) -> fun(st) { resume(st)(st) }
    case Put(st,resume) -> fun(_) { resume(())(st) }
  };
  f(st0)
}

fun example6() {
  runState(0, fun() {
    var x = 3;
    put(x);
    assert(x == get(), "Put;Get");
    var y = get();
    var z = get();
    assert(y == z, "Get;Get");
    put(x+1);
    put(x+2);
    assert(get() == x + 2, "Put;Put")
  })
}

##
## State: file i/o
##
typename FilePtr    = Option(FileDescr);

typename INode      = (loc:Int,lno:Int);
typename IList      = [(Int, INode)];  # INode index -> INode
typename Directory  = [(String, Int)]; # Filename -> INode index
typename DataRegion = [(Int, File)];   # Loc -> File

typename FileSystem = (dir:Directory,dregion:DataRegion,inodes:IList
                                    ,lnext:Int         ,inext:Int    );

sig fsys0 : FileSystem
var fsys0 = (dir = [("stdout", 0)], inodes = [(0, (loc=0, lno=1))], dregion = [(0, "")], lnext = 1, inext = 1);

sig fopen : (String, FileSystem) {Fail:Zero |_}~> (Int, FileSystem)
fun fopen(fname, fsys) { (lookup(fname, fsys.dir), fsys) }

sig ftruncate : (Int, FileSystem) {Fail:Zero |_}~> FileSystem
fun ftruncate(ino, fsys) {
  var inode = lookup(ino, fsys.inodes);
  var dregion = modify(inode.loc, "", fsys.dregion);
  (fsys with =dregion)
}

sig fcreate : (String, FileSystem) {Fail:Zero |_}~> (Int, FileSystem)
fun fcreate(fname, fsys) {
  if (has(fname, fsys.dir)) {
    var (ino, fsys) = fopen(fname, fsys);
    (ino, ftruncate(ino, fsys))
  } else {
    var loc = fsys.lnext;
    var dregion = (loc, "") :: fsys.dregion;

    var ino   = fsys.inext;
    var inode = (loc=loc,lno=1);
    var inodes = (ino, inode) :: fsys.inodes;

    var dir = (fname, ino) :: fsys.dir;
    (ino, (=dir, =dregion, =inodes, lnext=loc+1, inext=ino+1))
  }
}

sig ftruncate : (Int, FileSystem) {Fail:Zero |_}~> FileSystem
fun ftruncate(ino, fsys) {
  var inode = lookup(ino, fsys.inodes);
  var dregion = modify(inode.loc, "", fsys.dregion);
  (fsys with =dregion)
}

sig fopen : (String, FileSystem) {Fail:Zero |_}~> Int
fun fopen(fname, fsys) { lookup(fname, fsys.dir) }

sig fclose : (Int, FileSystem) ~> FileSystem
fun fclose(_, fsys) { fsys }

sig fwrite : (Int, String, FileSystem) {Fail:Zero |_}~> FileSystem
fun fwrite(ino, cs, fsys) {
  var inode = lookup(ino, fsys.inodes);
  var file = lookup(inode.loc, fsys.dregion);
  var file' = file ^^ cs;
  (fsys with dregion = modify(inode.loc, file', fsys.dregion))
}

sig fread : (Int, FileSystem) {Fail:Zero |_}~> String
fun fread(ino, fsys) {
  var inode = lookup(ino, fsys.inodes);
  lookup(inode.loc, fsys.dregion)
}

sig flink : (String, String, FileSystem) {Fail:Zero |_}~> FileSystem
fun flink(src, dest, fsys) {
  var ino    = lookup(dest, fsys.dir);
  var inode  = lookup(ino, fsys.inodes);
  var inode'  = (inode with lno = inode.lno + 1);
  var inodes = modify(ino, inode', fsys.inodes);

  var dir    = (src, ino) :: fsys.dir;
  (fsys with inodes = inodes, dir = dir)
}

sig funlink : (String, FileSystem) {Fail:Zero |_}~> FileSystem
fun funlink(fname, fsys) {
  var i = lookup(fname, fsys.dir);
  var dir = remove(fname, fsys.dir);

  var inode  = lookup(i, fsys.inodes);
  var inode'  = (inode with lno = inode.lno - 1);

  if (inode'.lno > 0) {
    var inodes = modify(i, inode', fsys.inodes);
    (fsys with inodes = inodes, dir = dir)
  } else {
    var dregion = remove(inode'.loc, fsys.dregion);
    var inodes = remove(i, fsys.inodes);
    (fsys with inodes = inodes, dir = dir, dregion = dregion)
  }
}

sig create : (String) {Create:(String) -> FileDescr |_}-> FileDescr
fun create(fname) { do Create(fname) }

sig truncate : (FileDescr) {Truncate:(FileDescr) -> () |_}-> ()
fun truncate(fd) { do Truncate(fd) }

sig open' : (String) {Open:(String) -> Option(FileDescr) |_}-> Option(FileDescr)
fun open'(fname) { do Open(fname) }

sig close : (FileDescr) {Close:(FileDescr) -> () |_}-> ()
fun close(fd) { do Close(fd) }

sig write : (FileDescr, String) {Write:(FileDescr, String) -> () |_}-> ()
fun write(fd, cs) { do Write(fd, cs) }

sig read : (FileDescr) {Read:(FileDescr) -> Option(String) |_}-> Option(String)
fun read(fd) { do Read(fd) }

sig link : (String, String) {Link:(String, String) -> () |_}-> ()
fun link(src, dest) { do Link(src, dest) }

sig unlink : (String) {Unlink:(String) -> () |_}-> ()
fun unlink(fname) { do Unlink(fname) }

sig injectState : (() { |e}~> a) -> () {Get:s,Put:(s) -> () |e}~> a
fun injectState(f) { (f : (() {Get:s,Put:(s) -> () |e}~> a) <- (() { |e}~> a)) }

sig fileIO : (Comp(a, { #Close:(FileDescr) -> ()
                        Create:(String) -> FileDescr
                      , Read:(FileDescr) -> Option(String)
                      , Open:(String) -> Option(FileDescr)
                      , Truncate:(FileDescr) -> ()
                      , Write:(FileDescr, String) -> ()
                      , Fail{p}|e}))
              {Create{_},Read{_},Open{_},Truncate{_},Write{_},Get:FileSystem,Put:(FileSystem) -> (),Fail{p} |e}~> a
fun fileIO(m) {
  handle(injectState(m)()) {
    case Return(x) -> x
    case Create(fname, resume) ->
      var ino = withDefault(-1, fun() {
        var (ino, fsys) = fcreate(fname, get());
        put(fsys); ino
      }); resume(ino)
    case Read(ino, resume) ->
      var contents = optionalise(fun() { fread(ino, get()) });
      resume(contents)
    case Open(fname, resume) ->
      var ino = optionalise(fun() { fopen(fname, get()) });
      resume(ino)
    case Truncate(ino, resume) ->
      withDefault((), fun() {
        var fsys = ftruncate(ino, get());
        put(fsys)
      }); resume(())
    case Write(ino, cs, resume) ->
      withDefault((), fun() {
        var fsys = fwrite(ino, cs, get());
        put(fsys)
      }); resume(())
  }
}

fun init(fsys, main) {
  runState(fsys, fun() {
    fileIO(fun() {
      timeshare(fun() {
        sessionmgr(Root, fun() {
          status(fun() {
            if(fork()) exit(0) else main()
          })
        })
      })
    })
  })
}

###
### Stream redirection
###
sig >> : (Comp(a, { Create: (String) -> FileDescr
                  , Write:(FileDescr, String) -> () |e}), String)
             { Create:(String) -> FileDescr
             , Write:(FileDescr, String) -> ()      |e}~> a
op f >> fname {
  var fd = create(fname);
  handle(f()) {
    case Return(x) -> x
    case Write(_, cs, resume) ->
      resume(write(fd, cs))
  }
}

sig >>> : (Comp(a, { Create: (String) -> FileDescr
                   , Open: (String)   -> Option (FileDescr)
                   , Write:(FileDescr, String) -> () |e}), String)
             { Create:(String) -> FileDescr
             , Open:(String) -> Option (FileDescr)
             , Write:(FileDescr, String) -> ()      |e}~> ()
op f >>> fname {
  var fd = switch (open'(fname)) {
      case None -> create(fname)
      case Some(fd) -> fd
  };
  handle(f()) {
    case Return(_) -> ()
    case Write(_, cs, resume) ->
      resume(write(fd, cs))
  }
}

fun example7() {
  if (fork()) {
    su(Alice);
    ritchie >> "ritchie.txt"
  } else {
    su(Bob);
    hamlet >> "hamlet"
  }
}

###
### TCP threeway handshake
###

sig strsplit : (Char, String) ~> [String]
fun strsplit(c, str) {
  fun loop(c, str, i, j) {
    if (i >= strlen(str)) [strsub(str, j, i - j)]
    else if (charAt(str, i) == c)
      strsub(str, j, i - j) :: loop(c, str, i+1, i+1)
    else
      loop(c,str, i+1, j)
  }
  loop(c, str, 0, 0)
}

sig read1 : (FileDescr) {Read:(FileDescr) -> Option(String) |e}-> Option(String)
fun read1(fd) {
  switch (read(fd)) {
    case Some("") -> None
    case x    -> x
  }
}

sig truncread : (FileDescr) {Read:(FileDescr) -> Option(String),Truncate:(FileDescr) -> () |_}-> Option(String)
fun truncread(fd) {
  var cs = read1(fd);
  truncate(fd); cs
}

sig synced : (Comp(Option(a), {Interrupt:() |e})) {Interrupt:() |e}~> a
fun synced(f) {
  switch (f()) {
    case None -> interrupt(); synced(f)
    case Some(x)  -> x
  }
}

sig fail : () {Fail:Zero |_}-> a
fun fail() { switch(do Fail) {} }

fun tcpclient(seq, inp, out) {
  write(out, "SYN " ^^ intToString(seq));
  var resp = synced(fun() { truncread(inp) });
  var [syn, ack] = strsplit(';', resp);
  var seqn = stringToInt(strsub(syn, 4, strlen(syn) - 4));
  var ackn = stringToInt(strsub(ack, 4, strlen(ack) - 4));
  if (ackn <> seq + 1) fail()
  else write(out, "ACK " ^^ intToString(seqn + 1))
}

fun tcpserver(seq, inp, out) {
  var req = synced(fun() { truncread(inp) });
  var reqn = stringToInt(strsub(req, 4, strlen(req) - 4));
  var resp = "SYN " ^^ intToString(seq) ^^ ";ACK " ^^ intToString(reqn + 1);
  write(out, resp);
  var resp' = synced(fun() { truncread(inp) });
  var ackn = stringToInt(strsub(resp', 4, strlen(resp') - 4));
  if (ackn <> seq + 1) fail()
  else ()
}

fun performTCP(tcpf, seq, inp, out) {
  var fd = create(whoami() ^^ ".log");
  handle(tcpf(seq, inp, out)) {
    case Return(_) -> write(fd, "Handshake completed.")
    case Fail(_) -> write(fd, "Handshake failed.")
   }
}

fun tcphandshake() {
  var (cfd, sfd) = (create("client.sock"), create("server.sock"));
  if (fork()) {
    su(Alice);
    performTCP(tcpclient, 42, cfd, sfd)
  } else {
    su(Bob);
    performTCP(tcpserver, 84, sfd, cfd)
  }
}

fun tcphandshakeFail() {
  var (cfd, sfd) = (create("client.sock"), create("server.sock"));
  if (fork()) {
    su(Alice);
    handle(performTCP(tcpclient, 42, cfd, sfd)) {
      case Write(fd, cs, resume) ->
        resume(if (strsub(cs, 0, 3) == "ACK") write(fd, "ACK 0")
               else write(fd, cs))
    }
  } else {
    su(Bob);
    performTCP(tcpserver, 84, sfd, cfd)
  }
}

#
# Grep
#
#sig grep : (String) {Await:Char,Yield:(Char) -> () |_}~> ()
fun grep(str) {
  var cs = explode(str);
  fun match(c,cs) {
    switch (cs) {
      case c' :: cs' ->
        if (c == '\n') fail()
        if (c == c') 
    }
  }
}