CPS IO -> Async Interaction Tree Effect Runtime

I'm deeply satisfied to be building an interaction tree runtime where
the interaction trees are themselves computed via and represented by
trees. It's trees all the way down.

README.md

@@ -2,13 +2,15 @@
 
 ## Introduction
 
-tricu (pronounced "tree-shoe") is a programming language experiment in Haskell. It is fundamentally based on the application of [Triage Calculus](https://olydis.medium.com/a-visual-introduction-to-tree-calculus-2f4a34ceffc2), an extended form of [Tree Calculus](https://github.com/barry-jay-personal/typed_tree_calculus/blob/main/typed_program_analysis.pdf), but minimal syntax sugar is included.
+tricu (pronounced "tree-shoe") is an experimental programming language written in Haskell. It is fundamentally based on the application of [Triage Calculus](https://olydis.medium.com/a-visual-introduction-to-tree-calculus-2f4a34ceffc2), an extended form of [Tree Calculus](https://github.com/barry-jay-personal/typed_tree_calculus/blob/main/typed_program_analysis.pdf). I will refer to this "family" of calculi as TC.
 
 tricu is the word for "tree" in Lojban: `(x1) is a tree of species/cultivar (x2)`.
 
+In the `ext/` directory there are implementations of TC evaluators and tooling in other languages. Here be dragons; beware.
+
 I have fully embraced the slopmachine (LLM-assisted development) for this project. Nothing is stable or sacred. We will discover sanity at the end of the journey but we won't strive for it until then.
 
-This README.md is human written. No other .md file will be until stabilization.
+This README.md is 100% human written. No other .md file will be until stabilization.
 
 ## Acknowledgements 
 

demos/monadicIO.tri

@@ -0,0 +1,86 @@
+!import "../lib/base.tri"      !Local
+!import "../lib/list.tri"      !Local
+!import "../lib/io.tri"        !Local
+
+-- Monadic IO in tricu
+--
+-- The IO system is a free monad interpreted by the host.  Primitive actions
+-- (putStr, readFile, writeFile, ...) do not carry their own continuations.
+-- Sequencing is performed by the single generic `bind` constructor.
+--
+--   pure x            -- lift a pure value into IO
+--   bind action k     -- run action, then apply k to its result
+--   thenIO a b        -- run a, discard its result, then run b
+--   mapIO action f    -- run action, then apply f to its result inside pure
+--
+-- File operations return a Result tree (see lib/base.tri):
+--   ok value     -- pair true  (pair value t)
+--   err code     -- pair false (pair code  t)
+--
+-- Use onReadFile / onWriteFile for convenient branching.
+
+-- ----------------------------------------------------------------------------
+-- Example 1: Greet and return a pure value.
+--   putStrLn writes to stdout; pure lifts "done" into IO.
+-- ----------------------------------------------------------------------------
+
+greet = (name :
+  bind (putStrLn (append "Hello, " name))
+    (_ : pure "done"))
+
+-- ----------------------------------------------------------------------------
+-- Example 2: Read a file safely.
+--   readFile returns a Result.  matchResult branches on ok / err.
+-- ----------------------------------------------------------------------------
+
+safeRead = (path :
+  bind (readFile path)
+    (result :
+      matchResult
+        (err rest : pure "missing")
+        (contents rest : pure contents)
+        result))
+
+-- ----------------------------------------------------------------------------
+-- Example 3: Write, then read back.
+--   thenIO discards the writeFile Result and continues.
+-- ----------------------------------------------------------------------------
+
+writeThenRead = (path text :
+  thenIO
+    (writeFile path text)
+    (readFile path))
+
+-- ----------------------------------------------------------------------------
+-- Example 4: Transform an IO result.
+--   mapIO applies a pure function to the value produced by an action.
+-- ----------------------------------------------------------------------------
+
+shout = (path :
+  mapIO (safeRead path)
+    (text : append text "!!!"))
+
+-- ----------------------------------------------------------------------------
+-- Example 5: Cooperative async.
+--   fork runs an action in the background.
+--   sleep suspends the current task for N milliseconds.
+--   await waits for a forked task and returns its value.
+--
+--   Here the child sleeps for 2 s while the parent prints immediately.
+--   The parent's message appears first, proving interleaving.
+-- ----------------------------------------------------------------------------
+
+asyncDemo = (
+  bind (fork
+    (bind (sleep 2000) (_ :
+    bind (putStrLn "Done sleeping!") (_ :
+    pure "child done"))))
+  (handle :
+  bind (putStrLn "Parent first!") (_ :
+  await handle)))
+
+-- ----------------------------------------------------------------------------
+-- Main action - run the async demo.
+-- ----------------------------------------------------------------------------
+
+main = io asyncDemo

lib/io.tri

@@ -3,80 +3,82 @@
 !import "conversions.tri" !Local
 
 -- IO constructors for host-interpreted interaction trees.
--- See docs/io-in-tricu.md for the full protocol.
+-- Free-monad style: Bind is the single sequencing mechanism.
 
 version = 1
 
 io        = action : pair "tricuIO" (pair version action)
-pure      = x      : pair 0          x
+
-putStr    = s k    : pair 1         (pair s k)
+pure      = x          : pair 0 x
-getLine   = k      : pair 2          k
+bind      = action k   : pair 1 (pair action k)
-readFile  = p k    : pair 3         (pair p k)
+
-writeFile = p c k  : pair 4         (pair p (pair c k))
+putStr    = s          : pair 10 s
+getLine   = pair 11 t
+
+readFile  = p          : pair 20 p
+writeFile = p c        : pair 21 (pair p c)
+
+ask       = pair 30 t
+local     = f action   : pair 31 (pair f action)
+
+get       = pair 40 t
+put       = s          : pair 41 s
+
+fork      = action     : pair 60 action
+await     = handle     : pair 61 handle
+yield     = pair 62 t
+sleep     = ms         : pair 63 ms
 
 -- ---------------------------------------------------------------------------
--- CPS sequencing helpers
+-- Generic sequencing combinators
 -- ---------------------------------------------------------------------------
 
--- Print a string and finish successfully.
+thenIO = a b : bind a (_ : b)
-print = s : putStr s (_ : pure t)
+mapIO  = action f : bind action (x : pure (f x))
 
--- Print a string plus newline and finish successfully.
+-- ---------------------------------------------------------------------------
-printLn = s : putStr (append s "\n") (_ : pure t)
+-- Convenience helpers
+-- ---------------------------------------------------------------------------
 
--- CPS print with newline.
+print    = s : bind (putStr s)       (_ : pure t)
-putStrLn = s k : putStr (append s "\n") k
+putStrLn = s : bind (putStr (append s "\n")) (_ : pure t)
-
--- Sequence after putStr, ignoring Unit.
-afterPutStr = s next : putStr s (_ : next)
-
--- Sequence after putStrLn, ignoring Unit.
-afterPutStrLn = s next : putStr (append s "\n") (_ : next)
 
 -- ---------------------------------------------------------------------------
 -- Result-aware file helpers
 -- ---------------------------------------------------------------------------
 
--- Read a file, forcing the caller to handle both success and error.
 onReadFile = (path errCase okCase :
-  readFile path (result :
+  bind (readFile path) (result :
     matchResult errCase okCase result))
 
--- Write a file, forcing the caller to handle both success and error.
 onWriteFile = (path contents errCase okCase :
-  writeFile path contents (result :
+  bind (writeFile path contents) (result :
     matchResult errCase okCase result))
 
 -- ---------------------------------------------------------------------------
 -- Convenience helpers for the common cases
 -- ---------------------------------------------------------------------------
 
--- Read a file; on error print a message and finish.
 readFileOrPrintError = (path okCase :
   onReadFile path
-    (err rest :
+    (err rest : putStrLn "Read failed")
-      putStrLn "Read failed" (_ :
-      pure t))
     okCase)
 
--- Write a file; on error print a message and finish.
 writeFileOrPrintError = (path contents okCase :
   onWriteFile path contents
-    (err rest :
+    (err rest : putStrLn "Write failed")
-      putStrLn "Write failed" (_ :
-      pure t))
     okCase)
 
--- Copy src to dst, then continue with k on success.
+copyFile = (src dst :
-copyFile = (src dst k :
+  bind (readFile src)
-  onReadFile src
+    (result :
-    (err rest :
+      matchResult
-      putStrLn "Read failed" (_ :
+        (err rest : putStrLn "Read failed")
-      pure t))
+        (contents rest :
-    (contents rest :
+          bind (writeFile dst contents)
-      onWriteFile dst contents
+            (wr :
-        (err rest :
+              matchResult
-          putStrLn "Write failed" (_ :
+                (err rest : putStrLn "Write failed")
-          pure t))
+                (ok rest   : pure t)
-        (ok rest :
+                wr))
-          k t)))
+        result))

notes/iodriver-updates.md

@@ -0,0 +1,749 @@
+Below is the implementation handoff for replacing the current recursive/rebuilding IO small-step interpreter with an explicit machine stack, primarily to support `Reader` via `ask` and `local`, while setting up the right shape for eventual async.
+
+## Goal
+
+Refactor `IODriver` from this model:
+
+```haskell
+stepIO :: IOPermissions -> T -> IO Step
+
+data Step
+  = Done T
+  | Continue T
+```
+
+to an explicit abstract machine:
+
+```haskell
+Machine = Runtime + current action + continuation frames
+```
+
+This is required because `local` is dynamically scoped. It needs to modify the Reader environment for a sub-computation, then restore the previous environment exactly when that sub-computation completes. The current “rebuild `Bind left' k`” approach has nowhere to store that restoration behavior.
+
+This change should support:
+
+```tricu
+ask
+local
+```
+
+now, and keep the structure compatible with future async suspension/resumption.
+
+Do not implement async in this pass.
+
+---
+
+## New action tags
+
+Extend the tricu IO action language with Reader tags:
+
+```tricu
+ask   = _        : pair 6 t
+local = f action : pair 7 (pair f action)
+```
+
+Host-side:
+
+```haskell
+data Action
+  = APure T
+  | ABind T T
+  | APutStr T
+  | AGetLine
+  | AReadFile T
+  | AWriteFile T T
+  | AAsk
+  | ALocal T T
+  deriving (Show)
+```
+
+Recommended tag allocation:
+
+```text
+0 = pure
+1 = bind
+2 = putStr
+3 = getLine
+4 = readFile
+5 = writeFile
+6 = ask
+7 = local
+```
+
+State tags can come later:
+
+```text
+8 = get
+9 = put
+```
+
+Do not add `bindR`, `bindS`, or `bindRS` yet. Reader is being added as an effect inside the existing IO action language, so the existing IO `bind` remains the only sequencing operator.
+
+---
+
+## New runtime model
+
+Add a runtime context:
+
+```haskell
+data Runtime = Runtime
+  { rtPerms :: IOPermissions
+  , rtEnv   :: T
+  }
+  deriving (Show)
+```
+
+Later this can become:
+
+```haskell
+data Runtime = Runtime
+  { rtPerms :: IOPermissions
+  , rtEnv   :: T
+  , rtState :: T
+  }
+```
+
+but for this pass, keep it minimal unless State is implemented at the same time.
+
+Add continuation frames:
+
+```haskell
+data Frame
+  = BindFrame T
+  | LocalFrame T
+  deriving (Show)
+```
+
+Frame meanings:
+
+```text
+BindFrame k:
+  When the current action produces value x, continue with apply k x.
+
+LocalFrame oldEnv:
+  When the current action produces value x, restore oldEnv, then continue with x.
+```
+
+Add the machine state:
+
+```haskell
+data Machine = Machine
+  { machineRuntime :: Runtime
+  , machineCurrent :: T
+  , machineFrames  :: [Frame]
+  }
+  deriving (Show)
+```
+
+Frames should be treated as a stack, with the head as the top:
+
+```haskell
+push frame machine = machine { machineFrames = frame : machineFrames machine }
+```
+
+---
+
+## New step result
+
+Replace the current `Step` with machine-oriented stepping:
+
+```haskell
+data Step
+  = Halt Runtime T
+  | Continue Machine
+  deriving (Show)
+```
+
+`Halt runtime value` means the entire IO program is done.
+
+`Continue machine` means the machine can take another step.
+
+---
+
+## Core stepping semantics
+
+The central function should become:
+
+```haskell
+stepMachine :: Machine -> IO Step
+```
+
+It should decode `machineCurrent`.
+
+### `pure`
+
+When the current action is `APure value`, do not immediately halt. First inspect the frame stack.
+
+Pseudo-code:
+
+```haskell
+finishValue :: Machine -> T -> IO Step
+finishValue machine value =
+  case machineFrames machine of
+    [] ->
+      pure (Halt (machineRuntime machine) value)
+
+    BindFrame k : rest ->
+      pure (Continue machine
+        { machineCurrent = apply k value
+        , machineFrames = rest
+        })
+
+    LocalFrame oldEnv : rest ->
+      let runtime' = (machineRuntime machine) { rtEnv = oldEnv }
+      in pure (Continue machine
+        { machineRuntime = runtime'
+        , machineCurrent = pureAction value
+        , machineFrames = rest
+        })
+```
+
+You will need a helper:
+
+```haskell
+pureAction :: T -> T
+pureAction x = Fork (ofNumber 0) x
+```
+
+This is important: restoring a `LocalFrame` should not discard the value. It restores the environment and re-enters the machine as `pure value`, allowing the next frame to receive the value.
+
+### `bind`
+
+For:
+
+```haskell
+ABind left k
+```
+
+do not recursively step `left`, and do not rebuild `Bind left' k`.
+
+Instead:
+
+```haskell
+Continue machine
+  { machineCurrent = left
+  , machineFrames  = BindFrame k : machineFrames machine
+  }
+```
+
+This is the major refactor. Continuations move out of the tree and into the frame stack.
+
+### `ask`
+
+For:
+
+```haskell
+AAsk
+```
+
+produce the current Reader environment:
+
+```haskell
+finishValue machine (rtEnv (machineRuntime machine))
+```
+
+or equivalently:
+
+```haskell
+Continue machine { machineCurrent = pureAction currentEnv }
+```
+
+Prefer `finishValue` because it avoids an extra step.
+
+### `local`
+
+For:
+
+```haskell
+ALocal f action
+```
+
+do:
+
+```haskell
+let runtime = machineRuntime machine
+    oldEnv  = rtEnv runtime
+    newEnv  = apply f oldEnv
+    runtime' = runtime { rtEnv = newEnv }
+
+Continue machine
+  { machineRuntime = runtime'
+  , machineCurrent = action
+  , machineFrames  = LocalFrame oldEnv : machineFrames machine
+  }
+```
+
+This is the central correctness point.
+
+`local` enters a scoped environment by pushing a restoration frame. When the scoped action finishes, `LocalFrame oldEnv` restores the previous environment and passes the produced value onward.
+
+Nested `local` works naturally because frames stack:
+
+```tricu
+local f (
+  local g ask
+)
+```
+
+becomes:
+
+```text
+push LocalFrame env0
+set env = f env0
+
+push LocalFrame env1
+set env = g env1
+
+ask returns env2
+
+pop LocalFrame env1
+restore env1
+
+pop LocalFrame env0
+restore env0
+```
+
+### Normal IO actions
+
+For host IO actions, perform the side effect and then call `finishValue`.
+
+Examples:
+
+```haskell
+APutStr str ->
+  case decodeString str "PutStr" of
+    Right s -> do
+      putStr s
+      finishValue machine Leaf
+    Left _ ->
+      finishValue machine (errResult 6)
+```
+
+```haskell
+AReadFile path ->
+  case decodeString path "ReadFile" of
+    Right p -> do
+      result <- ...
+      finishValue machine result
+    Left _ ->
+      finishValue machine (errResult 6)
+```
+
+Important: IO actions should no longer return `Done value` directly. They should return a value to the frame stack via `finishValue`.
+
+---
+
+## Decode changes
+
+Extend `decodeAction`:
+
+```haskell
+decodeAction :: T -> Either String Action
+decodeAction tree =
+  case tree of
+    Fork tag payload ->
+      case toNumber tag of
+        Right 0 -> Right (APure payload)
+
+        Right 1 -> case payload of
+          Fork left k -> Right (ABind left k)
+          _ -> Left "Invalid Bind: expected pair action continuation"
+
+        Right 2 -> Right (APutStr payload)
+
+        Right 3 -> Right AGetLine
+
+        Right 4 -> Right (AReadFile payload)
+
+        Right 5 -> case payload of
+          Fork path contents -> Right (AWriteFile path contents)
+          _ -> Left "Invalid WriteFile: expected pair path contents"
+
+        Right 6 -> Right AAsk
+
+        Right 7 -> case payload of
+          Fork f action -> Right (ALocal f action)
+          _ -> Left "Invalid Local: expected pair function action"
+
+        Right n -> Left $ "Unknown IO action tag: " ++ show n
+
+        Left err -> Left $ "Invalid action tag: " ++ err
+
+    _ ->
+      Left $ "Invalid action tree: expected pair tag payload, got " ++ show tree
+```
+
+---
+
+## Runner API
+
+Add a new Reader-aware runner:
+
+```haskell
+runIOWithEnv :: IOPermissions -> T -> T -> IO T
+runIOWithEnv perms env action = loop initialMachine
+  where
+    initialMachine = Machine
+      { machineRuntime = Runtime
+          { rtPerms = perms
+          , rtEnv = env
+          }
+      , machineCurrent = action
+      , machineFrames = []
+      }
+
+    loop machine = do
+      step <- stepMachine machine
+      case step of
+        Halt _ value       -> pure value
+        Continue machine'  -> loop machine'
+```
+
+Keep the existing API as a compatibility wrapper:
+
+```haskell
+runIO :: IOPermissions -> T -> IO T
+runIO perms action =
+  runIOWithEnv perms Leaf action
+```
+
+If State is added in the same branch, prefer:
+
+```haskell
+runIOWith :: IOPermissions -> T -> T -> T -> IO (T, T)
+```
+
+where:
+
+```text
+permissions
+initial reader env
+initial state
+action
+```
+
+returns:
+
+```text
+final result
+final state
+```
+
+But if this handoff is only for Reader, use `runIOWithEnv`.
+
+---
+
+## Permission helpers
+
+The current permission helper functions can mostly stay as-is, but they should read permissions from runtime:
+
+```haskell
+let perms = rtPerms (machineRuntime machine)
+```
+
+The current helpers are nested inside `stepIO`. After the refactor, either:
+
+1. keep them in a `where` block under `stepMachine`, or
+2. lift them to top-level helper functions.
+
+Prefer lifting pure/reusable helpers to top-level if this file is getting large:
+
+```haskell
+decodeString :: T -> String -> Either String String
+canonicalizeSafe :: FilePath -> IO (Either String FilePath)
+pathAllowed :: FilePath -> [FilePath] -> IO Bool
+tryReadFile :: FilePath -> IO T
+tryWriteFile :: FilePath -> String -> IO T
+okResult :: T -> T
+errResult :: Integer -> T
+ioErrorCode :: IOException -> Integer
+```
+
+This will make `stepMachine` much easier to read.
+
+---
+
+## `io.tri` changes
+
+Add the Reader constructors:
+
+```tricu
+ask   = _        : pair 6 t
+local = f action : pair 7 (pair f action)
+```
+
+No new bind is required.
+
+Example usage:
+
+```tricu
+program =
+  bind ask (env :
+    putStrLn env)
+```
+
+Example `local` usage:
+
+```tricu
+program =
+  bind ask (outer :
+  bind (local (env : append env "-inner")
+        (bind ask (inner :
+          pure inner)))
+    (result :
+      bind ask (after :
+        pure result)))
+```
+
+Expected behavior:
+
+```text
+outer ask sees original env
+inner ask sees transformed env
+after ask sees original env again
+```
+
+---
+
+## Tests to add
+
+Add tests around behavior, not implementation details.
+
+### 1. `ask` returns initial environment
+
+Program:
+
+```tricu
+io (bind ask (env : pure env))
+```
+
+Run with env:
+
+```text
+"dev"
+```
+
+Expected result:
+
+```text
+"dev"
+```
+
+### 2. `local` transforms environment
+
+Program:
+
+```tricu
+io (
+  local (env : append env "-local")
+    (bind ask (env : pure env))
+)
+```
+
+Initial env:
+
+```text
+"root"
+```
+
+Expected result:
+
+```text
+"root-local"
+```
+
+### 3. `local` restores environment afterward
+
+Program structure:
+
+```tricu
+bind ask (before :
+bind (local transform scopedAsk) (inside :
+bind ask (after :
+pure (pair before (pair inside after)))))
+```
+
+Initial env:
+
+```text
+"root"
+```
+
+Expected:
+
+```text
+pair "root" (pair "root-local" "root")
+```
+
+### 4. nested `local` composes correctly
+
+Program:
+
+```tricu
+local f (
+  local g ask
+)
+```
+
+Initial env:
+
+```text
+"root"
+```
+
+Example:
+
+```tricu
+f = x : append x "-f"
+g = x : append x "-g"
+```
+
+Expected inner ask:
+
+```text
+"root-f-g"
+```
+
+Also verify after both locals, environment is restored by doing a final `ask`.
+
+### 5. `local` result passes through bind correctly
+
+Program:
+
+```tricu
+bind
+  (local transform (pure "value"))
+  (x : pure x)
+```
+
+Expected:
+
+```text
+"value"
+```
+
+This catches a common bug where `LocalFrame` restores env but loses the value.
+
+### 6. IO still works through bind
+
+Existing IO tests should continue passing unchanged through `runIO`.
+
+### 7. IO inside local
+
+Program:
+
+```tricu
+local transform (
+  bind ask (env :
+  bind (putStrLn env) (_ :
+  pure env))
+)
+```
+
+Expected:
+
+```text
+prints transformed env
+returns transformed env
+```
+
+Then optionally ask after local to verify restoration.
+
+---
+
+## Invariants to preserve
+
+The implementation should maintain these invariants:
+
+```text
+1. The current action is always the next instruction to evaluate.
+
+2. The frame stack contains all pending continuations and cleanup scopes.
+
+3. Bind does not recursively step its left side.
+   It pushes BindFrame and switches current to the left action.
+
+4. local does not run its action to completion.
+   It pushes LocalFrame and switches current to the scoped action.
+
+5. Only LocalFrame restores Reader environment.
+
+6. State, when added later, should not be restored by LocalFrame.
+
+7. Existing runIO behavior remains source-compatible.
+```
+
+---
+
+## Common failure modes
+
+The likely bugs are:
+
+```text
+Bug: local leaks environment.
+Cause: setting rtEnv but never restoring it.
+Fix: push LocalFrame oldEnv and restore in finishValue.
+
+Bug: local restores environment but loses result.
+Cause: popping LocalFrame and halting directly.
+Fix: after restoration, continue with pureAction value.
+
+Bug: bind continuations run under the wrong env.
+Cause: LocalFrame and BindFrame pop order is wrong.
+Fix: use stack head as top. Push LocalFrame when entering local; push BindFrame when entering bind. Pop exactly one frame when a value is produced.
+
+Bug: existing IO bind tests fail.
+Cause: IO actions halt instead of passing result to finishValue.
+Fix: every completed primitive action should call finishValue.
+
+Bug: nested binds still rebuild trees.
+Cause: old ABind logic left in place.
+Fix: ABind should only push BindFrame and switch current to left.
+```
+
+---
+
+## Async relevance, but not implementation
+
+This machine representation is intentionally compatible with async.
+
+A future scheduler can store:
+
+```haskell
+Machine runtime current frames
+```
+
+when a task blocks, then resume it later by restoring the same `Machine`.
+
+Do not implement any of these now:
+
+```haskell
+AFork
+AAwait
+ASleep
+TaskId
+Scheduler
+Runnable queue
+Blocked table
+```
+
+But avoid designs that would make future suspension impossible, especially recursive “run sub-computation to completion” implementations of `local`. The point of the frame machine is that every effect remains small-step and resumable.
+
+---
+
+## Recommended implementation order
+
+1. Add `Runtime`, `Frame`, and `Machine`.
+2. Add `pureAction`.
+3. Replace `Step` with `Halt Runtime T | Continue Machine`.
+4. Implement `finishValue`.
+5. Rewrite `ABind` to push `BindFrame`.
+6. Rewrite existing primitive IO actions to call `finishValue`.
+7. Add `AAsk` and `ALocal`.
+8. Add `runIOWithEnv`.
+9. Rewrite `runIO` as a wrapper.
+10. Add `ask` and `local` to `io.tri`.
+11. Add Reader behavior tests.
+12. Run all existing IO tests and confirm no regressions.
+
+The key handoff instruction is: implement `local` as a continuation frame, not as a recursive nested run. This keeps the interpreter genuinely small-step and gives the eventual async runtime the exact representation it will need for suspension and resumption.

src/IODriver.hs

@@ -4,17 +4,28 @@ module IODriver
   , unsafePerms
   , checkIOSentinel
   , runIO
+  , runIOWithEnv
+  , runIOWith
   ) where
 
 import Research (T(..), apply, toString, toNumber, ofString, ofNumber)
 import System.IO (putStr, getLine)
 import qualified System.IO as IO
 import Control.Exception (try, IOException, SomeException)
-import System.Exit (die)
 import System.IO.Error (isDoesNotExistError, isPermissionError, isAlreadyExistsError)
 import Data.List (isPrefixOf)
 import System.FilePath (normalise, isRelative, (</>), addTrailingPathSeparator, splitDirectories)
 import System.Directory (canonicalizePath, doesPathExist, getCurrentDirectory)
+import qualified Data.Map.Strict as Map
+import Data.Map.Strict (Map)
+import qualified Data.Sequence as Seq
+import Data.Sequence (Seq, (|>), ViewL(..))
+import Data.Time.Clock (UTCTime, getCurrentTime, addUTCTime, diffUTCTime)
+import Control.Concurrent (threadDelay)
+
+-- ---------------------------------------------------------------------------
+-- Permissions
+-- ---------------------------------------------------------------------------
 
 data IOPermissions = IOPermissions
   { allowRead     :: [FilePath]
@@ -22,6 +33,7 @@ data IOPermissions = IOPermissions
   , allowReadAll  :: Bool
   , allowWriteAll :: Bool
   }
+  deriving (Show)
 
 defaultPerms :: IOPermissions
 defaultPerms = IOPermissions [] [] False False
@@ -41,91 +53,390 @@ checkIOSentinel tree =
         _ -> Left "sentinel mismatch (expected \"tricuIO\")"
     _ -> Left "root is not an IO sentinel pair"
 
-runIO :: IOPermissions -> T -> IO T
+-- ---------------------------------------------------------------------------
-runIO perms actionTree = go actionTree
+-- Runtime, Frames, and Machine
+-- ---------------------------------------------------------------------------
+
+data Runtime = Runtime
+  { rtPerms :: IOPermissions
+  , rtEnv   :: T
+  , rtState :: T
+  }
+  deriving (Show)
+
+data Frame
+  = BindFrame T
+  | LocalFrame T
+  deriving (Show)
+
+data Machine = Machine
+  { machineRuntime :: Runtime
+  , machineCurrent :: T
+  , machineFrames  :: [Frame]
+  }
+  deriving (Show)
+
+-- ---------------------------------------------------------------------------
+-- Result convention
+-- ---------------------------------------------------------------------------
+-- Direct-return actions pass the raw value to the continuation:
+--   pure, bind, putStr, getLine, ask, local, get, put,
+--   fork, await, yield, sleep
+--
+-- Result-return actions wrap the outcome as an ok/err pair:
+--   ok val  = Fork (Stem Leaf) (Fork val Leaf)   -- (t t) val t
+--   err code = Fork Leaf (Fork code Leaf)        -- t code t
+--   readFile, writeFile
+--
+-- Runtime protocol errors are returned as direct values via errResult.
+
+-- Error code ranges:
+--   1-19   host IO / filesystem errors
+--   20-39  policy / permission errors
+--   40-59  protocol / decode / type errors
+--   60-79  async errors
+--   80-99  scheduler / runtime errors
+
+-- Host IO / filesystem errors (1-19)
+errDoesNotExist, errPermission, errAlreadyExists, errIOOther :: Integer
+errDoesNotExist  = 1
+errPermission    = 2
+errAlreadyExists = 3
+errIOOther       = 4
+
+-- Policy / permission errors (20-39)
+errPolicyDeny :: Integer
+errPolicyDeny = 20
+
+-- Protocol / decode / type errors (40-59)
+errInvalidAction, errInvalidString :: Integer
+errInvalidAction = 40
+errInvalidString = 41
+
+-- Async errors (60-79)
+errInvalidHandle, errSelfAwait, errInvalidSleep :: Integer
+errInvalidHandle = 60
+errSelfAwait     = 61
+errInvalidSleep  = 62
+
+-- Scheduler / runtime errors (80-99)
+errDeadlock :: Integer
+errDeadlock = 80
+
+ioErrorCode :: IOException -> Integer
+ioErrorCode e
+  | isDoesNotExistError  e = errDoesNotExist
+  | isPermissionError    e = errPermission
+  | isAlreadyExistsError e = errAlreadyExists
+  | otherwise              = errIOOther
+
+okResult :: T -> T
+okResult val = Fork (Stem Leaf) (Fork val Leaf)
+
+errResult :: Integer -> T
+errResult code = Fork Leaf (Fork (ofNumber code) Leaf)
+
+pureAction :: T -> T
+pureAction x = Fork (ofNumber 0) x
+
+invalidAsyncHandleResult :: T
+invalidAsyncHandleResult = errResult errInvalidHandle
+
+selfAwaitResult :: T
+selfAwaitResult = errResult errSelfAwait
+
+deadlockResult :: T
+deadlockResult = errResult errDeadlock
+
+invalidSleepResult :: T
+invalidSleepResult = errResult errInvalidSleep
+
+-- ---------------------------------------------------------------------------
+-- Task identity and handles
+-- ---------------------------------------------------------------------------
+
+newtype TaskId = TaskId Integer
+  deriving (Eq, Ord, Show)
+
+taskHandle :: TaskId -> T
+taskHandle (TaskId n) =
+  Fork (ofString "task") (ofNumber n)
+
+decodeTaskHandle :: T -> Either String TaskId
+decodeTaskHandle tree =
+  case tree of
+    Fork tag nTree -> do
+      tagString <- toString tag
+      if tagString == "task"
+        then TaskId <$> toNumber nTree
+        else Left "invalid task handle tag"
+    _ ->
+      Left "invalid task handle"
+
+-- ---------------------------------------------------------------------------
+-- Free-monad action AST
+-- ---------------------------------------------------------------------------
+
+data Action
+  = APure T
+  | ABind T T
+  | APutStr T
+  | AGetLine
+  | AReadFile T
+  | AWriteFile T T
+  | AAsk
+  | ALocal T T
+  | AGet
+  | APut T
+  | AFork T
+  | AAwait T
+  | AYield
+  | ASleep T
+  deriving (Show)
+
+-- ---------------------------------------------------------------------------
+-- Action tag constants
+-- ---------------------------------------------------------------------------
+
+tagPure, tagBind :: Integer
+tagPure = 0
+tagBind = 1
+
+tagPutStr, tagGetLine :: Integer
+tagPutStr = 10
+tagGetLine = 11
+
+tagReadFile, tagWriteFile :: Integer
+tagReadFile = 20
+tagWriteFile = 21
+
+tagAsk, tagLocal :: Integer
+tagAsk = 30
+tagLocal = 31
+
+tagGet, tagPut :: Integer
+tagGet = 40
+tagPut = 41
+
+tagFork, tagAwait, tagYield, tagSleep :: Integer
+tagFork = 60
+tagAwait = 61
+tagYield = 62
+tagSleep = 63
+
+data Step
+  = Halt Runtime T
+  | Continue Machine
+  | ForkRequested T Machine
+  | AwaitRequested TaskId Machine
+  | YieldRequested Machine
+  | SleepRequested Integer Machine
+  deriving (Show)
+
+decodeAction :: T -> Either String Action
+decodeAction tree =
+  case tree of
+    Fork tag payload ->
+      case toNumber tag of
+        Right n | n == tagPure ->
+          Right (APure payload)
+
+        Right n | n == tagBind ->
+          case payload of
+            Fork left k -> Right (ABind left k)
+            _ -> Left "Invalid Bind: expected pair action continuation"
+
+        Right n | n == tagPutStr ->
+          Right (APutStr payload)
+
+        Right n | n == tagGetLine ->
+          Right AGetLine
+
+        Right n | n == tagReadFile ->
+          Right (AReadFile payload)
+
+        Right n | n == tagWriteFile ->
+          case payload of
+            Fork path contents -> Right (AWriteFile path contents)
+            _ -> Left "Invalid WriteFile: expected pair path contents"
+
+        Right n | n == tagAsk ->
+          Right AAsk
+
+        Right n | n == tagLocal ->
+          case payload of
+            Fork f action -> Right (ALocal f action)
+            _ -> Left "Invalid Local: expected pair function action"
+
+        Right n | n == tagGet ->
+          Right AGet
+
+        Right n | n == tagPut ->
+          Right (APut payload)
+
+        Right n | n == tagFork ->
+          Right (AFork payload)
+
+        Right n | n == tagAwait ->
+          Right (AAwait payload)
+
+        Right n | n == tagYield ->
+          Right AYield
+
+        Right n | n == tagSleep ->
+          Right (ASleep payload)
+
+        Right n ->
+          Left $ "Unknown IO action tag: " ++ show n
+
+        Left err ->
+          Left $ "Invalid action tag: " ++ err
+
+    _ ->
+      Left $ "Invalid action tree: expected pair tag payload, got " ++ show tree
+
+-- ---------------------------------------------------------------------------
+-- Small-step IO machine
+-- ---------------------------------------------------------------------------
+
+finishValue :: Machine -> T -> IO Step
+finishValue machine value =
+  case machineFrames machine of
+    [] ->
+      pure (Halt (machineRuntime machine) value)
+
+    BindFrame k : rest ->
+      pure (Continue machine
+        { machineCurrent = apply k value
+        , machineFrames = rest
+        })
+
+    LocalFrame oldEnv : rest ->
+      let runtime' = (machineRuntime machine) { rtEnv = oldEnv }
+      in pure (Continue machine
+          { machineRuntime = runtime'
+          , machineCurrent = pureAction value
+          , machineFrames = rest
+          })
+
+stepMachine :: Machine -> IO Step
+stepMachine machine =
+  case decodeAction (machineCurrent machine) of
+    Right action -> dispatch action
+    Left _       -> finishValue machine (errResult errInvalidAction)
   where
-    go tree =
+    dispatch action = case action of
-      case tree of
+      APure val ->
-        Fork tag payload -> do
+        finishValue machine val
-          tagNum <- case toNumber tag of
-            Right n  -> return n
-            Left err -> die $ "Invalid IO action tag: " ++ err
-          dispatch tagNum payload
-        _ -> die $ "Invalid IO action tree: expected pair tag payload, got " ++ show tree
 
-    dispatch tagNum payload = case tagNum of
+      ABind left k ->
-      0 -> return payload  -- Pure
+        pure (Continue machine
+          { machineCurrent = left
+          , machineFrames = BindFrame k : machineFrames machine
+          })
 
-      1 -> case payload of
+      APutStr str ->
-        Fork str k -> do
+        case decodeString str "PutStr" of
-          s <- decodeString str "PutStr"
+          Right s -> do
-          putStr s
+            putStr s
-          go (apply k Leaf)
+            finishValue machine Leaf
-        _ -> die "Invalid PutStr payload: expected pair string continuation"
+          Left _ ->
+            finishValue machine (errResult errInvalidString)
 
-      2 -> do
+      AGetLine -> do
         line <- getLine
-        go (apply payload (ofString line))
+        finishValue machine (ofString line)
 
-      3 -> case payload of
+      AReadFile path ->
-        Fork path k -> do
+        case decodeString path "ReadFile" of
-          p <- decodeString path "ReadFile"
+          Right p -> do
-          mDeny <- checkReadPerm p
+            mDeny <- checkReadPerm p
-          case mDeny of
-            Just denied -> go (apply k denied)
-            Nothing -> do
-              content <- tryReadFile p
-              go (apply k content)
-        _ -> die "Invalid ReadFile payload: expected pair path continuation"
-
-      4 -> case payload of
-        Fork path rest -> case rest of
-          Fork contents k -> do
-            p <- decodeString path "WriteFile"
-            c <- decodeString contents "WriteFile"
-            mDeny <- checkWritePerm p
             case mDeny of
-              Just denied -> go (apply k denied)
+              Just denied -> finishValue machine denied
-              Nothing -> do
+              Nothing     -> tryReadFile p >>= finishValue machine
-                res <- tryWriteFile p c
+          Left _ -> finishValue machine (errResult errInvalidString)
-                go (apply k res)
-          _ -> die "Invalid WriteFile payload: expected pair contents continuation"
-        _ -> die "Invalid WriteFile payload: expected pair path (pair contents continuation)"
 
-      _ -> die $ "Unknown IO action tag: " ++ show tagNum
+      AWriteFile path contents ->
+        case decodeString path "WriteFile" of
+          Right p ->
+            case decodeString contents "WriteFile" of
+              Right c -> do
+                mDeny <- checkWritePerm p
+                case mDeny of
+                  Just denied -> finishValue machine denied
+                  Nothing     -> tryWriteFile p c >>= finishValue machine
+              Left _ -> finishValue machine (errResult errInvalidString)
+          Left _ -> finishValue machine (errResult errInvalidString)
 
-    decodeString t ctx =
+      AAsk ->
-      case toString t of
+        finishValue machine (rtEnv (machineRuntime machine))
-        Right s  -> return s
-        Left err -> die $ "Invalid " ++ ctx ++ " string: " ++ err
 
+      ALocal f action' ->
+        let runtime = machineRuntime machine
+            oldEnv  = rtEnv runtime
+            newEnv  = apply f oldEnv
+            runtime' = runtime { rtEnv = newEnv }
+        in pure (Continue machine
+            { machineRuntime = runtime'
+            , machineCurrent = action'
+            , machineFrames = LocalFrame oldEnv : machineFrames machine
+            })
+
+      AGet ->
+        finishValue machine (rtState (machineRuntime machine))
+
+      APut newState ->
+        let runtime' = (machineRuntime machine) { rtState = newState }
+        in finishValue (machine { machineRuntime = runtime' }) Leaf
+
+      AFork childAction ->
+        pure (ForkRequested childAction machine)
+
+      AAwait handleTree ->
+        case decodeTaskHandle handleTree of
+          Right taskId ->
+            pure (AwaitRequested taskId machine)
+          Left _ ->
+            finishValue machine invalidAsyncHandleResult
+
+      AYield ->
+        pure (YieldRequested machine)
+
+      ASleep msTree ->
+        case toNumber msTree of
+          Right ms | ms >= 0 ->
+            pure (SleepRequested ms machine)
+          _ ->
+            finishValue machine invalidSleepResult
+
+    -- Permission and IO helpers
     checkReadPerm p =
-      if allowReadAll perms
+      if allowReadAll (rtPerms (machineRuntime machine))
         then return Nothing
         else do
           mp <- canonicalizeSafe p
           case mp of
-            Left _    -> return $ Just policyErrResult
+            Left _     -> return $ Just policyErrResult
             Right path -> do
-              allowed <- pathAllowed path (allowRead perms)
+              allowed <- pathAllowed path (allowRead (rtPerms (machineRuntime machine)))
               if allowed
                 then return Nothing
                 else return $ Just policyErrResult
 
     checkWritePerm p =
-      if allowWriteAll perms
+      if allowWriteAll (rtPerms (machineRuntime machine))
         then return Nothing
         else do
           mp <- canonicalizeSafe p
           case mp of
-            Left _    -> return $ Just policyErrResult
+            Left _     -> return $ Just policyErrResult
             Right path -> do
-              allowed <- pathAllowed path (allowWrite perms)
+              allowed <- pathAllowed path (allowWrite (rtPerms (machineRuntime machine)))
               if allowed
                 then return Nothing
                 else return $ Just policyErrResult
 
-    policyErrResult = errResult 5
+    policyErrResult = errResult errPolicyDeny
 
     canonicalizeSafe :: FilePath -> IO (Either String FilePath)
     canonicalizeSafe p = do
@@ -187,12 +498,225 @@ runIO perms actionTree = go actionTree
         Right () -> return $ okResult Leaf
         Left e   -> return $ errResult (ioErrorCode e)
 
-    okResult val  = Fork (Stem Leaf) (Fork val Leaf)          -- pair true (pair val t)
+    decodeString t ctx =
-    errResult code = Fork Leaf (Fork (ofNumber code) Leaf)    -- pair false (pair code t)
+      case toString t of
+        Right s  -> Right s
+        Left _ -> Left $ "Invalid " ++ ctx ++ " string"
 
-    ioErrorCode :: IOException -> Integer
+-- ---------------------------------------------------------------------------
-    ioErrorCode e
+-- Scheduler
-      | isDoesNotExistError  e = 1
+-- ---------------------------------------------------------------------------
-      | isPermissionError    e = 2
+
-      | isAlreadyExistsError e = 3
+data TaskStatus
-      | otherwise              = 4
+  = Runnable Machine
+  | BlockedOn TaskId Machine
+  | Sleeping UTCTime Machine
+  | Completed Runtime T
+  deriving (Show)
+
+data Scheduler = Scheduler
+  { schedulerNextTaskId :: Integer
+  , schedulerRunnable   :: Seq TaskId
+  , schedulerTasks      :: Map TaskId TaskStatus
+  }
+  deriving (Show)
+
+initialScheduler :: Machine -> Scheduler
+initialScheduler mainMachine =
+  Scheduler
+    { schedulerNextTaskId = 1
+    , schedulerRunnable = Seq.singleton (TaskId 0)
+    , schedulerTasks = Map.singleton (TaskId 0) (Runnable mainMachine)
+    }
+
+runtimeOfStatus :: TaskStatus -> Runtime
+runtimeOfStatus (Runnable machine) = machineRuntime machine
+runtimeOfStatus (BlockedOn _ machine) = machineRuntime machine
+runtimeOfStatus (Sleeping _ machine) = machineRuntime machine
+runtimeOfStatus (Completed runtime _) = runtime
+
+wakeAwaiters :: TaskId -> T -> Scheduler -> Scheduler
+wakeAwaiters targetId value scheduler =
+  let (newlyRunnable, tasks') =
+        Map.mapAccumWithKey wakeOne [] (schedulerTasks scheduler)
+      queue' = foldl (|>) (schedulerRunnable scheduler) (reverse newlyRunnable)
+  in scheduler { schedulerTasks = tasks', schedulerRunnable = queue' }
+  where
+    wakeOne acc tid (BlockedOn blockedTarget machine)
+      | blockedTarget == targetId =
+          let machine' = machine { machineCurrent = pureAction value }
+          in (tid : acc, Runnable machine')
+    wakeOne acc _ status = (acc, status)
+
+wakeDueSleepers :: Scheduler -> IO Scheduler
+wakeDueSleepers scheduler = do
+  now <- getCurrentTime
+  let (newlyRunnable, tasks') =
+        Map.mapAccumWithKey (wakeOne now) [] (schedulerTasks scheduler)
+      queue' = foldl (|>) (schedulerRunnable scheduler) (reverse newlyRunnable)
+  pure scheduler { schedulerTasks = tasks', schedulerRunnable = queue' }
+  where
+    wakeOne now acc tid (Sleeping wakeTime machine)
+      | wakeTime <= now = (tid : acc, Runnable machine)
+    wakeOne _ acc _ status = (acc, status)
+
+nearestSleepTime :: Scheduler -> Maybe UTCTime
+nearestSleepTime = Map.foldl' minSleep Nothing . schedulerTasks
+  where
+    minSleep acc (Sleeping t _) = Just $ maybe t (min t) acc
+    minSleep acc _ = acc
+
+resumeCurrentWith :: TaskId -> T -> Machine -> Scheduler -> IO Scheduler
+resumeCurrentWith taskId value machine scheduler =
+  let machine' = machine { machineCurrent = pureAction value }
+  in pure scheduler
+      { schedulerTasks = Map.insert taskId (Runnable machine') (schedulerTasks scheduler)
+      , schedulerRunnable = schedulerRunnable scheduler |> taskId
+      }
+
+handleStep :: TaskId -> Step -> Scheduler -> IO Scheduler
+handleStep taskId (Continue machine) scheduler =
+  pure scheduler
+    { schedulerTasks = Map.insert taskId (Runnable machine) (schedulerTasks scheduler)
+    , schedulerRunnable = schedulerRunnable scheduler |> taskId
+    }
+
+handleStep taskId (Halt _runtime value) scheduler =
+  pure (wakeAwaiters taskId value scheduler')
+  where
+    scheduler' = scheduler
+      { schedulerTasks = Map.insert taskId (Completed _runtime value) (schedulerTasks scheduler)
+      }
+
+handleStep parentId (ForkRequested childAction parentMachine) scheduler =
+  let childId = TaskId (schedulerNextTaskId scheduler)
+      handle = taskHandle childId
+
+      parentMachine' =
+        parentMachine { machineCurrent = pureAction handle }
+
+      childMachine =
+        Machine
+          { machineRuntime = machineRuntime parentMachine
+          , machineCurrent = childAction
+          , machineFrames = []
+          }
+
+      tasks' =
+        Map.insert parentId (Runnable parentMachine') $
+        Map.insert childId (Runnable childMachine) $
+        schedulerTasks scheduler
+
+      queue' =
+        schedulerRunnable scheduler |> parentId |> childId
+
+  in pure scheduler
+      { schedulerNextTaskId = schedulerNextTaskId scheduler + 1
+      , schedulerTasks = tasks'
+      , schedulerRunnable = queue'
+      }
+
+handleStep currentId (AwaitRequested targetId machine) scheduler
+  | targetId == currentId =
+      resumeCurrentWith currentId selfAwaitResult machine scheduler
+
+  | otherwise =
+      case Map.lookup targetId (schedulerTasks scheduler) of
+        Nothing ->
+          resumeCurrentWith currentId invalidAsyncHandleResult machine scheduler
+
+        Just (Completed _ value) ->
+          resumeCurrentWith currentId value machine scheduler
+
+        Just _ ->
+          pure scheduler
+            { schedulerTasks =
+                Map.insert currentId (BlockedOn targetId machine) (schedulerTasks scheduler)
+            }
+
+handleStep taskId (YieldRequested machine) scheduler =
+  resumeCurrentWith taskId Leaf machine scheduler
+
+handleStep taskId (SleepRequested ms machine) scheduler = do
+  now <- getCurrentTime
+  let seconds = fromIntegral ms / 1000
+      wakeTime = addUTCTime seconds now
+      machine' = machine { machineCurrent = pureAction Leaf }
+  pure scheduler
+    { schedulerTasks =
+        Map.insert taskId (Sleeping wakeTime machine') (schedulerTasks scheduler)
+    }
+
+handleNoRunnable :: Scheduler -> IO Scheduler
+handleNoRunnable scheduler =
+  case nearestSleepTime scheduler of
+    Just wakeTime -> do
+      now <- getCurrentTime
+      let micros = max 0 (floor (diffUTCTime wakeTime now * 1000000))
+      threadDelay micros
+      wakeDueSleepers scheduler
+
+    Nothing ->
+      case Map.lookup (TaskId 0) (schedulerTasks scheduler) of
+        Just status ->
+          pure scheduler
+            { schedulerTasks =
+                Map.insert (TaskId 0)
+                  (Completed (runtimeOfStatus status) deadlockResult)
+                  (schedulerTasks scheduler)
+            }
+        Nothing ->
+          pure scheduler
+
+schedulerStep :: Scheduler -> IO Scheduler
+schedulerStep scheduler = do
+  scheduler1 <- wakeDueSleepers scheduler
+  case Seq.viewl (schedulerRunnable scheduler1) of
+    EmptyL ->
+      handleNoRunnable scheduler1
+
+    taskId :< restQueue ->
+      case Map.lookup taskId (schedulerTasks scheduler1) of
+        Just (Runnable machine) -> do
+          step <- stepMachine machine
+          handleStep taskId step scheduler1 { schedulerRunnable = restQueue }
+
+        _ ->
+          pure scheduler1 { schedulerRunnable = restQueue }
+
+runScheduler :: Scheduler -> IO (T, T)
+runScheduler scheduler =
+  case Map.lookup (TaskId 0) (schedulerTasks scheduler) of
+    Just (Completed runtime value) ->
+      pure (value, rtState runtime)
+
+    _ ->
+      schedulerStep scheduler >>= runScheduler
+
+-- ---------------------------------------------------------------------------
+-- Public API
+-- ---------------------------------------------------------------------------
+
+runIOWith :: IOPermissions -> T -> T -> T -> IO (T, T)
+runIOWith perms env initialState action =
+  runScheduler (initialScheduler initialMachine)
+  where
+    initialMachine = Machine
+      { machineRuntime = Runtime
+          { rtPerms = perms
+          , rtEnv = env
+          , rtState = initialState
+          }
+      , machineCurrent = action
+      , machineFrames = []
+      }
+
+runIOWithEnv :: IOPermissions -> T -> T -> IO T
+runIOWithEnv perms env action = do
+  (result, _) <- runIOWith perms env Leaf action
+  pure result
+
+runIO :: IOPermissions -> T -> IO T
+runIO perms action = do
+  (result, _) <- runIOWith perms Leaf Leaf action
+  pure result

test/Spec.hs

@@ -8,12 +8,13 @@ import REPL
 import Research
 import Wire
 import ContentStore
-import IODriver (IOPermissions(..), checkIOSentinel, runIO, unsafePerms)
+import IODriver (IOPermissions(..), checkIOSentinel, runIO, runIOWithEnv, runIOWith, unsafePerms, defaultPerms)
 
 import Control.Exception      (evaluate, try, SomeException)
 import Control.Monad         (forM_)
 import Control.Monad.IO.Class (liftIO)
 import System.IO.Temp         (withSystemTempDirectory)
+import System.Directory      (createDirectory)
 import Data.Bits              (xor)
 import Data.Char              (digitToInt)
 import Data.List              (isInfixOf)
@@ -1262,7 +1263,8 @@ byteListUtilities = testGroup "Byte List Utility Tests"
 
 ioDriverTests :: TestTree
 ioDriverTests = testGroup "IO driver tests"
-  [ testCase "readFile through onReadFile returns file contents" $
+  [ -- Existing behaviour tests
+    testCase "readFile through onReadFile returns file contents" $
       withSystemTempDirectory "tricu-io-read" $ \dir -> do
         let sourcePath = dir ++ "/input.txt"
         writeFile sourcePath "abc123"
@@ -1300,13 +1302,496 @@ ioDriverTests = testGroup "IO driver tests"
             , "    (second rest : pure (append first second))))"
             ]
         final @?= ofString "abcdef"
+
+    -- Monad law tests
+  , testCase "left identity: bind (pure x) f == f x" $ do
+      left <- runIOSource $
+        unlines
+          [ "f = x : pure (append x \"!\")"
+          , "main = io (bind (pure \"abc\") f)"
+          ]
+      right <- runIOSource $
+        unlines
+          [ "f = x : pure (append x \"!\")"
+          , "main = io (f \"abc\")"
+          ]
+      left @?= right
+      left @?= ofString "abc!"
+
+  , testCase "right identity: bind m pure == m" $
+      withSystemTempDirectory "tricu-io-right-id" $ \dir -> do
+        let path = dir ++ "/input.txt"
+        writeFile path "abc"
+        left <- runIOSource $
+          unlines
+            [ "main = io (bind (readFile \"" ++ path ++ "\")"
+            , "  (result : pure result))"
+            ]
+        right <- runIOSource $
+          unlines
+            [ "main = io (readFile \"" ++ path ++ "\")"
+            ]
+        left @?= right
+        left @?= ioOkResult (ofString "abc")
+
+  , testCase "associativity: bind (bind m f) g == bind m (x : bind (f x) g)" $
+      withSystemTempDirectory "tricu-io-assoc" $ \dir -> do
+        let path = dir ++ "/input.txt"
+        writeFile path "abc"
+        left <- runIOSource $
+          unlines
+            [ "m = readFile \"" ++ path ++ "\""
+            , "f = result : matchResult (err rest : pure \"read failed\") (contents rest : pure (append contents \"-f\")) result"
+            , "g = value : pure (append value \"-g\")"
+            , "main = io (bind (bind m f) g)"
+            ]
+        right <- runIOSource $
+          unlines
+            [ "m = readFile \"" ++ path ++ "\""
+            , "f = result : matchResult (err rest : pure \"read failed\") (contents rest : pure (append contents \"-f\")) result"
+            , "g = value : pure (append value \"-g\")"
+            , "main = io (bind m (x : bind (f x) g))"
+            ]
+        left @?= right
+        left @?= ofString "abc-f-g"
+
+  , testCase "associativity preserves error flow" $
+      withSystemTempDirectory "tricu-io-assoc-err" $ \dir -> do
+        let missingPath = dir ++ "/missing.txt"
+        left <- runIOSource $
+          unlines
+            [ "m = readFile \"" ++ missingPath ++ "\""
+            , "f = result : matchResult (err rest : pure \"handled\") (contents rest : pure (append contents \"-ok\")) result"
+            , "g = value : pure (append value \"-g\")"
+            , "main = io (bind (bind m f) g)"
+            ]
+        right <- runIOSource $
+          unlines
+            [ "m = readFile \"" ++ missingPath ++ "\""
+            , "f = result : matchResult (err rest : pure \"handled\") (contents rest : pure (append contents \"-ok\")) result"
+            , "g = value : pure (append value \"-g\")"
+            , "main = io (bind m (x : bind (f x) g))"
+            ]
+        left @?= right
+        left @?= ofString "handled-g"
+
+  , testCase "bind defers continuation until left action completes" $
+      withSystemTempDirectory "tricu-io-lazy-k" $ \dir -> do
+        let path = dir ++ "/created.txt"
+        final <- runIOSource $
+          unlines
+            [ "main = io (bind (writeFile \"" ++ path ++ "\" \"created\")"
+            , "  (_ : readFile \"" ++ path ++ "\"))"
+            ]
+        final @?= ioOkResult (ofString "created")
+
+    -- Primitive effect shape tests
+  , testCase "readFile without continuation returns Result" $
+      withSystemTempDirectory "tricu-io-raw-read" $ \dir -> do
+        let path = dir ++ "/input.txt"
+        writeFile path "abc"
+        final <- runIOSource $
+          unlines
+            [ "main = io (readFile \"" ++ path ++ "\")"
+            ]
+        final @?= ioOkResult (ofString "abc")
+
+  , testCase "writeFile then readFile executes exactly once" $
+      withSystemTempDirectory "tricu-io-once" $ \dir -> do
+        let path = dir ++ "/test.txt"
+        final <- runIOSource $
+          unlines
+            [ "main = io (bind (writeFile \"" ++ path ++ "\" \"abc\")"
+            , "  (_ : readFile \"" ++ path ++ "\"))"
+            ]
+        final @?= ioOkResult (ofString "abc")
+
+  , testCase "sequencing order is left-to-right" $
+      withSystemTempDirectory "tricu-io-order" $ \dir -> do
+        let path = dir ++ "/test.txt"
+        final <- runIOSource $
+          unlines
+            [ "main = io (bind (writeFile \"" ++ path ++ "\" \"a\")"
+            , "  (_ : bind (writeFile \"" ++ path ++ "\" \"ab\")"
+            , "    (_ : readFile \"" ++ path ++ "\")))"
+            ]
+        final @?= ioOkResult (ofString "ab")
+
+  , testCase "thenIO sequences two actions and discards first result" $
+      withSystemTempDirectory "tricu-io-then" $ \dir -> do
+        let path = dir ++ "/test.txt"
+        final <- runIOSource $
+          unlines
+            [ "main = io (thenIO (writeFile \"" ++ path ++ "\" \"x\")"
+            , "  (readFile \"" ++ path ++ "\"))"
+            ]
+        final @?= ioOkResult (ofString "x")
+
+  , testCase "bind does not short-circuit on readFile error" $
+      withSystemTempDirectory "tricu-io-no-short" $ \dir -> do
+        let path = dir ++ "/missing.txt"
+        final <- runIOSource $
+          unlines
+            [ "main = io (bind (readFile \"" ++ path ++ "\")"
+            , "  (result : pure \"continued\"))"
+            ]
+        final @?= ofString "continued"
+
+  , testCase "mapIO transforms pure value" $ do
+      final <- runIOSource $
+        unlines
+          [ "main = io (mapIO (pure \"abc\") (x : append x \"!\"))"
+          ]
+      final @?= ofString "abc!"
+
+    -- Malformed action tests
+  , testCase "unknown IO action tag returns err result" $ do
+      final <- runIOSource "main = io (pair 99 t)"
+      final @?= ioErrResult 40
+
+  , testCase "malformed Bind returns err result" $ do
+      final <- runIOSource "main = io (pair 1 t)"
+      final @?= ioErrResult 40
+
+  , testCase "malformed ReadFile payload returns err result" $ do
+      final <- runIOSource "main = io (readFile (t t))"
+      final @?= ioErrResult 41
+
+    -- Permission tests
+  , testCase "allowed read path succeeds" $
+      withSystemTempDirectory "tricu-io-allowed" $ \dir -> do
+        let path = dir ++ "/allowed.txt"
+        writeFile path "allowed"
+        let perms = defaultPerms { allowRead = [path] }
+        result <- runIOSourceWithPerms perms $
+          unlines
+            [ "main = io (readFile \"" ++ path ++ "\")"
+            ]
+        result @?= ioOkResult (ofString "allowed")
+
+  , testCase "readFile denied path returns err result" $
+      withSystemTempDirectory "tricu-io-read-denied" $ \dir -> do
+        let allowedPath = dir ++ "/allowed.txt"
+            deniedPath = dir ++ "/denied.txt"
+        writeFile allowedPath "allowed"
+        writeFile deniedPath "denied"
+        let perms = defaultPerms { allowRead = [allowedPath] }
+        result <- runIOSourceWithPerms perms $
+          unlines
+            [ "main = io (readFile \"" ++ deniedPath ++ "\")"
+            ]
+        result @?= ioErrResult 20
+
+  , testCase "writeFile denied path returns err result" $
+      withSystemTempDirectory "tricu-io-write-denied" $ \dir -> do
+        let allowedPath = dir ++ "/allowed.txt"
+            deniedPath = dir ++ "/denied.txt"
+        let perms = defaultPerms { allowWrite = [allowedPath] }
+        result <- runIOSourceWithPerms perms $
+          unlines
+            [ "main = io (writeFile \"" ++ deniedPath ++ "\" \"x\")"
+            ]
+        result @?= ioErrResult 20
+
+  , testCase "path prefix does not allow prefix bypass" $
+      withSystemTempDirectory "tricu-io-prefix" $ \dir -> do
+        let allowedDir = dir ++ "/foo"
+            bypassPath = dir ++ "/foobar/secret.txt"
+        createDirectory allowedDir
+        createDirectory (dir ++ "/foobar")
+        writeFile bypassPath "secret"
+        let perms = defaultPerms { allowRead = [allowedDir] }
+        result <- runIOSourceWithPerms perms $
+          unlines
+            [ "main = io (readFile \"" ++ bypassPath ++ "\")"
+            ]
+        result @?= ioErrResult 20
+
+    -- Pure test
+  , testCase "pure performs no effects" $ do
+      final <- runIOSource "main = io (pure \"abc\")"
+      final @?= ofString "abc"
+
+    -- Reader tests
+  , testCase "ask returns initial environment" $ do
+      final <- runIOSourceWithEnv unsafePerms (ofString "dev") $
+        unlines
+          [ "main = io (bind ask (env : pure env))"
+          ]
+      final @?= ofString "dev"
+
+  , testCase "local transforms environment" $ do
+      final <- runIOSourceWithEnv unsafePerms (ofString "root") $
+        unlines
+          [ "main = io (local (env : append env \"-local\") (bind ask (env : pure env)))"
+          ]
+      final @?= ofString "root-local"
+
+  , testCase "local restores environment afterward" $ do
+      final <- runIOSourceWithEnv unsafePerms (ofString "root") $
+        unlines
+          [ "main = io (bind ask (before :"
+          , "  bind (local (env : append env \"-local\") (bind ask (env : pure env))) (inside :"
+          , "  bind ask (after :"
+          , "  pure (pair before (pair inside after))))))"
+          ]
+      final @?= Fork (ofString "root") (Fork (ofString "root-local") (ofString "root"))
+
+  , testCase "nested local composes correctly" $ do
+      final <- runIOSourceWithEnv unsafePerms (ofString "root") $
+        unlines
+          [ "f = x : append x \"-f\""
+          , "g = x : append x \"-g\""
+          , "main = io (bind"
+          , "  (local f (local g (bind ask (env : pure env))))"
+          , "  (inner :"
+          , "  bind ask (after :"
+          , "  pure (pair inner after))))"
+          ]
+      final @?= Fork (ofString "root-f-g") (ofString "root")
+
+  , testCase "local result passes through bind correctly" $ do
+      final <- runIOSourceWithEnv unsafePerms (ofString "root") $
+        unlines
+          [ "main = io (bind"
+          , "  (local (env : append env \"-local\") (pure \"value\"))"
+          , "  (x : pure x))"
+          ]
+      final @?= ofString "value"
+
+  , testCase "IO inside local uses transformed environment and restores after" $ do
+      final <- runIOSourceWithEnv unsafePerms (ofString "root") $
+        unlines
+          [ "main = io (bind"
+          , "  (local (env : append env \"-local\")"
+          , "    (bind ask (env : pure env)))"
+          , "  (result :"
+          , "  bind ask (after :"
+          , "  pure (pair result after))))"
+          ]
+      final @?= Fork (ofString "root-local") (ofString "root")
+
+	, testCase "local does not affect outer bind continuation" $ do
+	    final <- runIOSourceWithEnv unsafePerms (ofString "root") $
+	      unlines
+	        [ "main = io (bind"
+	        , "  (local (env : append env \"-local\") (pure \"x\"))"
+	        , "  (_ : bind ask (env : pure env)))"
+	        ]
+	    final @?= ofString "root"
+
+	, testCase "local environment persists across inner binds" $ do
+	    final <- runIOSourceWithEnv unsafePerms (ofString "root") $
+	      unlines
+	        [ "main = io (local (env : append env \"-local\")"
+	        , "  (bind (pure t) (_ :"
+	        , "  bind ask (env : pure env))))"
+	        ]
+	    final @?= ofString "root-local"
+
+	, testCase "local restores environment when scoped action returns error value" $ do
+	    final <- runIOSourceWithEnv defaultPerms (ofString "root") $
+	      unlines
+	        [ "main = io (bind"
+	        , "  (local (env : append env \"-local\") (readFile \"definitely-missing.txt\"))"
+	        , "  (_ : bind ask (env : pure env)))"
+	        ]
+	    final @?= ofString "root"
+
+    -- State tests
+  , testCase "get returns initial state" $ do
+      (final, st) <- runIOSourceWith unsafePerms Leaf (ofNumber 42) $
+        unlines
+          [ "main = io (bind get (s : pure s))"
+          ]
+      final @?= ofNumber 42
+      st @?= ofNumber 42
+
+  , testCase "put updates state" $ do
+      (final, st) <- runIOSourceWith unsafePerms Leaf (ofNumber 0) $
+        unlines
+          [ "main = io (bind (put 100) (_ : bind get (s : pure s)))"
+          ]
+      final @?= ofNumber 100
+      st @?= ofNumber 100
+
+  , testCase "state persists through bind" $ do
+      (final, st) <- runIOSourceWith unsafePerms Leaf (ofNumber 5) $
+        unlines
+          [ "main = io (bind get (s1 :"
+          , "  bind (put (succ s1)) (_ :"
+          , "  bind get (s2 :"
+          , "  pure (pair s1 s2)))))"
+          ]
+      final @?= Fork (ofNumber 5) (ofNumber 6)
+      st @?= ofNumber 6
+
+  , testCase "local does not restore state" $ do
+      (final, st) <- runIOSourceWith unsafePerms Leaf (ofNumber 0) $
+        unlines
+          [ "main = io (bind (put 10) (_ :"
+          , "  bind (local (env : env) (put 20)) (_ :"
+          , "  bind get (s :"
+          , "  pure s))))"
+          ]
+      final @?= ofNumber 20
+      st @?= ofNumber 20
+
+  , testCase "state and reader are independent" $ do
+      (final, st) <- runIOSourceWith unsafePerms (ofString "hello") (ofNumber 42) $
+        unlines
+          [ "main = io (bind ask (env :"
+          , "  bind get (s :"
+          , "  pure (pair env s))))"
+          ]
+      final @?= Fork (ofString "hello") (ofNumber 42)
+      st @?= ofNumber 42
+
+    -- Async tests
+  , testCase "fork returns handle and await returns child value" $ do
+      (final, st) <- runIOSourceWith unsafePerms Leaf Leaf $
+        unlines
+          [ "main = io (bind (fork (pure \"child\")) (h :"
+          , "  await h))"
+          ]
+      final @?= ofString "child"
+      st @?= Leaf
+
+  , testCase "main completion abandons unawaited child" $ do
+      (final, _) <- runIOSourceWith unsafePerms Leaf Leaf $
+        unlines
+          [ "main = io (bind (fork (pure \"child\")) (_ :"
+          , "  pure \"main\"))"
+          ]
+      final @?= ofString "main"
+
+  , testCase "fork captures reader environment at fork point" $ do
+      (final, _) <- runIOSourceWith unsafePerms (ofString "root") Leaf $
+        unlines
+          [ "main = io (local (env : append env \"-local\")"
+          , "  (bind (fork (bind ask (env : pure env))) (h :"
+          , "  await h)))"
+          ]
+      final @?= ofString "root-local"
+
+  , testCase "fork inside local captures child env and parent restores env" $ do
+      (final, _) <- runIOSourceWith unsafePerms (ofString "root") Leaf $
+        unlines
+          [ "main = io (bind"
+          , "  (local (env : append env \"-local\")"
+          , "    (fork (bind ask (env : pure env))))"
+          , "  (h : bind ask (after :"
+          , "  bind (await h) (child :"
+          , "  pure (pair after child)))))"
+          ]
+      final @?= Fork (ofString "root") (ofString "root-local")
+
+  , testCase "fork copies state and child state does not merge" $ do
+      (final, st) <- runIOSourceWith unsafePerms Leaf (ofNumber 0) $
+        unlines
+          [ "main = io (bind (put 1) (_ :"
+          , "  bind (fork (bind (put 99) (_ : bind get (s : pure s)))) (h :"
+          , "  bind (put 2) (_ :"
+          , "  bind (await h) (childState :"
+          , "  bind get (parentState :"
+          , "  pure (pair childState parentState)))))))"
+          ]
+      final @?= Fork (ofNumber 99) (ofNumber 2)
+      st @?= ofNumber 2
+
+  , testCase "multiple awaiters receive same completed value" $ do
+      (final, _) <- runIOSourceWith unsafePerms Leaf Leaf $
+        unlines
+          [ "main = io (bind (fork (pure \"done\")) (h :"
+          , "  bind (await h) (a :"
+          , "  bind (await h) (b :"
+          , "  pure (pair a b)))))"
+          ]
+      final @?= Fork (ofString "done") (ofString "done")
+
+  , testCase "self await returns async error" $ do
+      (final, _) <- runIOSourceWith unsafePerms Leaf Leaf $
+        unlines
+          [ "main = io (await (pair \"task\" 0))"
+          ]
+      final @?= ioErrResult 61
+
+  , testCase "await invalid handle returns async error" $ do
+      (final, _) <- runIOSourceWith unsafePerms Leaf Leaf $
+        unlines
+          [ "main = io (await 123)"
+          ]
+      final @?= ioErrResult 60
+
+  , testCase "yield returns unit and resumes continuation" $ do
+      (final, _) <- runIOSourceWith unsafePerms Leaf Leaf $
+        unlines
+          [ "main = io (bind yield (_ : pure \"after\"))"
+          ]
+      final @?= ofString "after"
+
+  , testCase "sleep resumes continuation" $ do
+      (final, _) <- runIOSourceWith unsafePerms Leaf Leaf $
+        unlines
+          [ "main = io (bind (sleep 1) (_ : pure \"awake\"))"
+          ]
+      final @?= ofString "awake"
+
+  , testCase "await waits for sleeping child" $ do
+      (final, _) <- runIOSourceWith unsafePerms Leaf Leaf $
+        unlines
+          [ "main = io (bind (fork (bind (sleep 1) (_ : pure \"awake\"))) (h :"
+          , "  await h))"
+          ]
+      final @?= ofString "awake"
+
+	, testCase "await waits for sleeping child and returns child value" $ do
+		  (final, st) <- runIOSourceWith unsafePerms Leaf Leaf $
+		    unlines
+		      [ "main = io (bind (fork (bind (sleep 1) (_ : pure \"child done\"))) (h :"
+		      , "  await h))"
+		      ]
+		  final @?= ofString "child done"
+		  st @?= Leaf
+
+	, testCase "sleep inside bind resumes as unit" $ do
+		  (final, st) <- runIOSourceWith unsafePerms Leaf Leaf $
+		    unlines
+		      [ "main = io (bind (sleep 1) (_ : pure \"awake\"))"
+		      ]
+		  final @?= ofString "awake"
+		  st @?= Leaf
+
+  , testCase "fork await returns child value" $ do
+      (final, st) <- runIOSourceWith unsafePerms Leaf Leaf $
+        unlines
+          [ "main = io (bind (fork (pure \"child done\")) (h :"
+          , "  await h))"
+          ]
+      final @?= ofString "child done"
+      st @?= Leaf
   ]
 
-runIOSource :: String -> IO T
+runIOSourceWith :: IOPermissions -> T -> T -> String -> IO (T, T)
-runIOSource source = do
+runIOSourceWith perms readerEnv initialState source = do
   ioEnv <- evaluateFile "./lib/io.tri"
-  env <- evalTricuWithStore Nothing ioEnv (parseTricu source)
+  evalEnv <- evalTricuWithStore Nothing ioEnv (parseTricu source)
-  case checkIOSentinel (mainResult env) of
+  case checkIOSentinel (mainResult evalEnv) of
-    Right (1, action) -> runIO unsafePerms action
+    Right (1, action) -> runIOWith perms readerEnv initialState action
     Right (v, _)      -> assertFailure ("Unsupported IO ABI version: " ++ show v)
     Left err          -> assertFailure ("Expected IO sentinel: " ++ err)
+
+runIOSource :: String -> IO T
+runIOSource source = fmap fst $ runIOSourceWith unsafePerms Leaf Leaf source
+
+runIOSourceWithPerms :: IOPermissions -> String -> IO T
+runIOSourceWithPerms perms source = fmap fst $ runIOSourceWith perms Leaf Leaf source
+
+runIOSourceWithEnv :: IOPermissions -> T -> String -> IO T
+runIOSourceWithEnv perms readerEnv source = fmap fst $ runIOSourceWith perms readerEnv Leaf source
+
+ioOkResult :: T -> T
+ioOkResult val = Fork (Stem Leaf) (Fork val Leaf)
+
+ioErrResult :: Integer -> T
+ioErrResult code = Fork Leaf (Fork (ofNumber code) Leaf)