243: Env Comonad (CoReader)

Difficulty: 5 Level: Master Pair a value with a read-only environment — the exact dual of the Reader monad.

The Problem This Solves

You're writing an expression evaluator, a configuration-aware computation, or a rendering pipeline where every node needs access to some global context (a variable table, a config struct, a render context). The Reader monad lets you thread that environment through function calls. The Env comonad does the dual thing: it carries the environment alongside a value, and you can run context-sensitive transformations on it. The classic use case is an expression tree evaluator. Instead of passing a `HashMap<String, i64>` as a parameter through every recursive call, you wrap the expression with the environment: `Env { env: variable_table, value: expression }`. Then `eval` is just a function from `&Env<VarEnv, Expr>` to `i64` — a comonad `extend` operation. The Env comonad also models the "staged computation" pattern: compute a value in one environment, then re-run it in a different environment using `local`. This is how configuration layers work in real systems.

The Intuition

A comonad is the categorical dual of a monad. Where a monad wraps values and lets you chain computations into the context (`bind`), a comonad wraps values and lets you extract them out of context (`extract`) and extend computations over the context (`extend`). The Env comonad is `Env<E, A> = (E, A)` — just a pair. The environment `E` is the context; `A` is the current value.

`extract`: get the value, ignore the environment. `(env, value) → value`.
`ask`: get the environment, ignore the value. `(env, value) → env`.
`extend f`: apply `f` to the whole pair to produce a new value, keeping the same environment. `(env, f(env, value))`. The environment is never modified by extend.
`local f`: run with a temporarily modified environment. Like a scoped config override.

The key duality with Reader: `Reader<E, A>` is `E → A` (a function that consumes an environment to produce a value). `Env<E, A>` is `(E, A)` (a value that carries its environment with it). One computes with the environment; the other is annotated by it. The three comonad laws are mirror images of monad laws: 1. `extract(extend f x) = f(x)` — extending then extracting gets back the original function result 2. `extend extract = identity` — extending with `extract` does nothing 3. `extend f . extend g = extend (f . extend g)` — associativity of extend

How It Works in Rust

pub struct Env<E, A> {
 pub env: E,
 pub value: A,
}

impl<E: Clone, A: Clone> Env<E, A> {
 // The two fundamental operations:
 pub fn extract(&self) -> A { self.value.clone() }
 pub fn ask(&self) -> E    { self.env.clone() }

 // extend: compute a new value using the full (env, value) pair.
 // The environment is passed through unchanged.
 pub fn extend<B: Clone>(&self, f: impl Fn(&Env<E, A>) -> B) -> Env<E, B> {
     Env {
         env: self.env.clone(),    // environment unchanged
         value: f(self),           // new value from the computation
     }
 }

 // local: temporarily modify the environment.
 pub fn local(&self, f: impl Fn(E) -> E) -> Env<E, A> {
     Env { env: f(self.env.clone()), value: self.value.clone() }
 }
}

Expression evaluation using `extend`:

// Wrap expression + variable table in Env
let node = Env::new(variable_table, expression);

// eval is the function we pass to extend:
// it can read both the environment (variable table) and value (expression)
let result = node.extend(eval);  // eval: &Env<VarEnv, Expr> -> i64

// To run in a different environment (e.g., staging vs production):
let staging = node.local(|mut env| { env.insert("x".into(), 99); env });
let staging_result = staging.extend(eval);

The expression evaluator recursively builds new `Env` nodes with the same environment, threading it down without explicit parameter passing:

Expr::Add(l, r) => {
 eval(&Env::new(env_expr.env.clone(), *l.clone()))
 + eval(&Env::new(env_expr.env.clone(), *r.clone()))
}

What This Unlocks

Expression evaluators and interpreters: thread a variable binding context without explicit parameter threading.
Configuration layers: `local` models "run this computation with config overrides" — like a stack-based config system.
Context-aware rendering: annotate each tree node with its rendering context, transform with `extend`, read with `extract`.

Key Differences

Concept	OCaml	Rust
Env type	`type ('e, 'a) env = { env: 'e; value: 'a }`	`struct Env<E, A> { env: E, value: A }`
extract	Pattern match or `.value` field	`self.value.clone()` (Clone bound required)
extend	Higher-kinded function via modules/functors	Generic method with `impl Fn`
local	Close over env value	Close over env value, same pattern
Comonad typeclass	`type class Comonad w` via module signature	Trait, or just implement methods directly

/// Env Comonad (CoReader): pairs a value with a read-only environment.
///
/// Env<E, A> = (E, A)
///
///   ask:     get the environment
///   extract: get the value (ignoring the environment)
///   extend:  run a computation that can read the environment, replacing the value
///
/// This is the exact dual of the Reader monad:
///   Reader<E, A>: E -> A     (environment produces a value)
///   Env<E, A>:   (E, A)     (environment stored alongside a value)
///
/// Use case: expression evaluation with a variable binding environment.

use std::collections::HashMap;

/// Env<E, A>: a value A tagged with environment E.
#[derive(Debug, Clone)]
pub struct Env<E, A> {
    pub env: E,
    pub value: A,
}

impl<E: Clone, A: Clone> Env<E, A> {
    pub fn new(env: E, value: A) -> Self {
        Env { env, value }
    }

    /// Comonad: extract = get the value (ignore the environment).
    pub fn extract(&self) -> A {
        self.value.clone()
    }

    /// ask: get the environment.
    pub fn ask(&self) -> E {
        self.env.clone()
    }

    /// Comonad: extend.
    /// Run `f` with access to the full Env, replace the value.
    /// The environment is unchanged.
    pub fn extend<B: Clone>(&self, f: impl Fn(&Env<E, A>) -> B) -> Env<E, B> {
        Env {
            env: self.env.clone(),
            value: f(self),
        }
    }

    /// duplicate: Env<E, A> -> Env<E, Env<E, A>>
    pub fn duplicate(&self) -> Env<E, Env<E, A>> {
        Env {
            env: self.env.clone(),
            value: self.clone(),
        }
    }

    /// local: run with a modified environment.
    pub fn local(&self, f: impl Fn(E) -> E) -> Env<E, A> {
        Env {
            env: f(self.env.clone()),
            value: self.value.clone(),
        }
    }
}

// ── Expression Evaluator via Env Comonad ─────────────────────────────────────

type VarEnv = HashMap<String, i64>;

/// Simple arithmetic expression AST.
#[derive(Debug, Clone)]
pub enum Expr {
    Lit(i64),
    Var(String),
    Add(Box<Expr>, Box<Expr>),
    Mul(Box<Expr>, Box<Expr>),
    Let(String, Box<Expr>, Box<Expr>),  // let x = e1 in e2
}

/// Evaluate an expression within an environment.
/// Uses `Env` comonad: the environment is the binding context.
fn eval(env_expr: &Env<VarEnv, Expr>) -> i64 {
    match &env_expr.value {
        Expr::Lit(n) => *n,
        Expr::Var(x) => *env_expr.env.get(x).unwrap_or(&0),
        Expr::Add(l, r) => {
            let env = env_expr.env.clone();
            eval(&Env::new(env.clone(), *l.clone()))
            + eval(&Env::new(env, *r.clone()))
        }
        Expr::Mul(l, r) => {
            let env = env_expr.env.clone();
            eval(&Env::new(env.clone(), *l.clone()))
            * eval(&Env::new(env, *r.clone()))
        }
        Expr::Let(x, e1, body) => {
            let val = eval(&Env::new(env_expr.env.clone(), *e1.clone()));
            let mut new_env = env_expr.env.clone();
            new_env.insert(x.clone(), val);
            eval(&Env::new(new_env, *body.clone()))
        }
    }
}

fn main() {
    println!("=== Env Comonad (CoReader) ===\n");
    println!("Env<E, A> = (E, A)");
    println!("Dual of Reader monad: environment carried alongside value.\n");

    // Basic Env comonad operations
    let e: Env<String, i32> = Env::new("production".to_string(), 42);
    println!("Env(\"production\", 42):");
    println!("  extract = {}", e.extract());
    println!("  ask     = \"{}\"", e.ask());

    // extend: compute based on both env and value
    let e2 = e.extend(|env| {
        if env.ask() == "production" {
            env.extract() * 100
        } else {
            env.extract()
        }
    });
    println!("  extend (if prod: *100): {}", e2.extract());

    // local: modify the environment
    let e3 = e.local(|_| "staging".to_string());
    let e4 = e3.extend(|env| {
        if env.ask() == "production" {
            env.extract() * 100
        } else {
            env.extract() + 1
        }
    });
    println!("  local(staging) then extend: {}", e4.extract());

    // Comonad law: extract . extend f = f
    let f = |env: &Env<String, i32>| env.extract() + env.ask().len() as i32;
    let extended = e.extend(f);
    assert_eq!(extended.extract(), f(&e), "Law 1: extract . extend f = f");
    println!("\nComonad law 1: extract . extend f = f: {} = {} ✓",
        extended.extract(), f(&e));

    // Expression evaluator
    println!("\n=== Expression Evaluator ===\n");
    let mut vars = HashMap::new();
    vars.insert("x".to_string(), 10_i64);
    vars.insert("y".to_string(), 3_i64);
    vars.insert("z".to_string(), 7_i64);

    // x + 2*y = 16
    let e1 = Expr::Add(
        Box::new(Expr::Var("x".to_string())),
        Box::new(Expr::Mul(Box::new(Expr::Lit(2)), Box::new(Expr::Var("y".to_string())))),
    );
    let result1 = eval(&Env::new(vars.clone(), e1));
    println!("x + 2*y = {} (expected 16)", result1);

    // (x + z) * y = 51
    let e2 = Expr::Mul(
        Box::new(Expr::Add(
            Box::new(Expr::Var("x".to_string())),
            Box::new(Expr::Var("z".to_string())),
        )),
        Box::new(Expr::Var("y".to_string())),
    );
    let result2 = eval(&Env::new(vars.clone(), e2));
    println!("(x+z)*y = {} (expected 51)", result2);

    // let a = x + 1 in a * y
    let e3 = Expr::Let(
        "a".to_string(),
        Box::new(Expr::Add(Box::new(Expr::Var("x".to_string())), Box::new(Expr::Lit(1)))),
        Box::new(Expr::Mul(Box::new(Expr::Var("a".to_string())), Box::new(Expr::Var("y".to_string())))),
    );
    let result3 = eval(&Env::new(vars.clone(), e3));
    println!("let a = x+1 in a*y = {} (expected 33)", result3);

    // Extend: evaluate and also return the env size
    let env_node = Env::new(vars.clone(), Expr::Var("x".to_string()));
    let eval_and_env_size = env_node.extend(|w| (eval(w), w.env.len()));
    println!("\nEval + env_size: val={}, env_vars={}", eval_and_env_size.extract().0, eval_and_env_size.extract().1);
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_extract_ask() {
        let e = Env::new("env".to_string(), 42_i32);
        assert_eq!(e.extract(), 42);
        assert_eq!(e.ask(), "env");
    }

    #[test]
    fn test_extend() {
        let e = Env::new(10_i32, 5_i32);
        let e2 = e.extend(|w| w.extract() + w.ask());
        assert_eq!(e2.extract(), 15);
        assert_eq!(e2.ask(), 10); // environment unchanged
    }

    #[test]
    fn test_local() {
        let e = Env::new(1_i32, 100_i32);
        let e2 = e.local(|n| n * 2);
        assert_eq!(e2.ask(), 2);
        assert_eq!(e2.extract(), 100);
    }

    #[test]
    fn test_eval_lit() {
        let vars = HashMap::new();
        let result = eval(&Env::new(vars, Expr::Lit(42)));
        assert_eq!(result, 42);
    }

    #[test]
    fn test_eval_add() {
        let mut vars = HashMap::new();
        vars.insert("a".to_string(), 3_i64);
        vars.insert("b".to_string(), 4_i64);
        let expr = Expr::Add(
            Box::new(Expr::Var("a".to_string())),
            Box::new(Expr::Var("b".to_string())),
        );
        assert_eq!(eval(&Env::new(vars, expr)), 7);
    }

    #[test]
    fn test_comonad_law1() {
        // extract . extend f = f
        let e = Env::new(5_i32, 10_i32);
        let f = |w: &Env<i32, i32>| w.extract() * w.ask();
        assert_eq!(e.extend(f).extract(), f(&e));
    }
}

(* Env comonad (also called CoReader): pairs a value with a read-only environment.
   Dual of the Reader monad.
   extract: get the value (ignoring environment)
   extend: access the environment while computing *)

type ('e, 'a) env = Env of 'e * 'a

let ask     (Env (e, _)) = e
let extract (Env (_, a)) = a

let extend (Env (e, a)) f =
  Env (e, f (Env (e, a)))

(* A simple expression evaluator with environment *)
type expr =
  | Lit of int
  | Var of string
  | Add of expr * expr
  | Mul of expr * expr

type env_map = (string * int) list

let rec eval_expr (Env (env, expr)) =
  match expr with
  | Lit n        -> n
  | Var x        -> List.assoc x env
  | Add (l, r)   ->
    eval_expr (Env (env, l)) + eval_expr (Env (env, r))
  | Mul (l, r)   ->
    eval_expr (Env (env, l)) * eval_expr (Env (env, r))

let () =
  let env = [("x", 10); ("y", 3); ("z", 7)] in

  let e1 = Add (Var "x", Mul (Lit 2, Var "y")) in (* x + 2*y = 16 *)
  let result1 = eval_expr (Env (env, e1)) in
  Printf.printf "x + 2*y = %d\n" result1;

  let e2 = Mul (Add (Var "x", Var "z"), Var "y") in (* (x+z)*y = 51 *)
  let result2 = eval_expr (Env (env, e2)) in
  Printf.printf "(x+z)*y = %d\n" result2;

  (* Extend: derive new computation from environment *)
  let w = Env (env, "x") in
  let w' = extend w (fun (Env (e, k)) -> List.assoc k e * 2) in
  Printf.printf "x * 2 via extend = %d\n" (extract w')