242: Store Comonad

/// Store Comonad: models a mutable "store" (array/function at an index).
///
/// Store s a = (s -> a, s)
///
///   extract: apply the getter at the current position
///   extend:  at each position s', compute f(Store(getter, s'))
///   peek:    look at any position without moving the focus
///   pos:     get the current focus position
///
/// The Store comonad is the dual of the State monad.
/// A classic application: 1D cellular automata, where each cell's
/// next state depends on its neighborhood.

/// Store<S, A>: a "getter" function `S -> A` plus a focus position `S`.
#[derive(Clone)]
pub struct Store<S, A> {
    getter: std::rc::Rc<dyn Fn(S) -> A>,
    position: S,
}

impl<S: Clone + 'static, A: Clone + 'static> Store<S, A> {
    pub fn new(getter: impl Fn(S) -> A + 'static, position: S) -> Self {
        Store {
            getter: std::rc::Rc::new(getter),
            position,
        }
    }

    /// Comonad: extract = apply the getter at the current position.
    pub fn extract(&self) -> A {
        (self.getter)(self.position.clone())
    }

    /// peek: look up any position (not just the current one).
    pub fn peek(&self, s: S) -> A {
        (self.getter)(s)
    }

    /// pos: get the current focus position.
    pub fn pos(&self) -> S {
        self.position.clone()
    }

    /// seeks: move the focus to a new position.
    pub fn seek(&self, s: S) -> Store<S, A> {
        Store {
            getter: self.getter.clone(),
            position: s,
        }
    }

    /// Comonad: extend.
    /// Builds a new Store where each position s' holds f(Store(getter, s')).
    pub fn extend<B: Clone + 'static>(
        &self,
        f: impl Fn(&Store<S, A>) -> B + 'static,
    ) -> Store<S, B> {
        let original = self.clone();
        let f = std::rc::Rc::new(f);
        Store {
            getter: std::rc::Rc::new(move |s: S| {
                let moved = original.seek(s);
                f(&moved)
            }),
            position: self.position.clone(),
        }
    }

    /// duplicate: Store s a -> Store s (Store s a)
    /// Each position now holds the Store focused there.
    pub fn duplicate(&self) -> Store<S, Store<S, A>> {
        let original = self.clone();
        Store {
            getter: std::rc::Rc::new(move |s: S| original.seek(s)),
            position: self.position.clone(),
        }
    }
}

// ── 1D Cellular Automaton via Store Comonad ───────────────────────────────────

/// Build a Store from a fixed-size array with wrap-around indexing.
fn make_cell_store(cells: &[bool]) -> Store<i32, bool> {
    let cells = cells.to_vec();
    let len = cells.len() as i32;
    Store::new(
        move |i: i32| {
            let idx = ((i % len) + len) % len;
            cells[idx as usize]
        },
        0,
    )
}

/// Rule 30: one step of the cellular automaton.
/// next(i) = f(left, center, right)
fn rule30_step(store: &Store<i32, bool>) -> bool {
    let i = store.pos();
    let left = store.peek(i - 1);
    let center = store.extract();
    let right = store.peek(i + 1);
    match (left, center, right) {
        (true,  true,  true)  => false,
        (true,  true,  false) => false,
        (true,  false, true)  => false,
        (true,  false, false) => true,
        (false, true,  true)  => true,
        (false, true,  false) => true,
        (false, false, true)  => true,
        (false, false, false) => false,
    }
}

/// Extract a generation array from a Store.
fn extract_gen(store: &Store<i32, bool>, size: usize) -> Vec<bool> {
    (0..size as i32).map(|i| store.peek(i)).collect()
}

/// Print one generation.
fn print_gen(gen: &[bool]) {
    for &b in gen {
        print!("{}", if b { '#' } else { '.' });
    }
    println!();
}

fn main() {
    println!("=== Store Comonad ===\n");
    println!("Store s a = (s -> a, s)");
    println!("extract: read at current position");
    println!("extend:  transform the getter, keeping the store structure\n");

    // Basic Store operations
    let store = Store::new(|i: i32| i * i, 4_i32);
    println!("Store (square function) at position 4:");
    println!("  extract = {}", store.extract());    // 16
    println!("  peek(3) = {}", store.peek(3));      // 9
    println!("  peek(7) = {}", store.peek(7));      // 49

    // seek: move focus
    let moved = store.seek(5);
    println!("  After seek(5): extract = {}", moved.extract()); // 25

    // extend: add 1 to every position's value
    let incremented = store.extend(|s| s.extract() + 1);
    println!("\nAfter extend(+1):");
    println!("  pos=4 -> {}", incremented.peek(4)); // 17
    println!("  pos=3 -> {}", incremented.peek(3)); // 10

    // Comonad law: extract . extend f = f
    let f = |s: &Store<i32, i32>| s.extract() * 2;
    let extended = store.extend(f);
    assert_eq!(extended.extract(), f(&store), "Law 1: extract . extend f = f");
    println!("\nComonad law 1 (extract . extend f = f): {} = {} ✓",
        extended.extract(), f(&store));

    // Cellular automaton
    println!("\n=== Rule 30 Cellular Automaton via Store Comonad ===\n");
    let size = 21_usize;
    let mut cells = vec![false; size];
    cells[size / 2] = true; // single cell in the center

    print_gen(&cells);

    // Each generation = extend rule30_step over the Store
    let mut s = make_cell_store(&cells);
    for _ in 0..5 {
        s = s.extend(rule30_step);
        print_gen(&extract_gen(&s, size));
    }

    println!("\nStore comonad models the cellular automaton elegantly:");
    println!("  extend rule30_step   applies the rule to every cell simultaneously.");
    println!("  peek(i±1)            accesses neighbors without mutation.");
}

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

    #[test]
    fn test_extract() {
        let s = Store::new(|i: i32| i * 3, 4_i32);
        assert_eq!(s.extract(), 12);
    }

    #[test]
    fn test_peek() {
        let s = Store::new(|i: i32| i + 10, 0_i32);
        assert_eq!(s.peek(5), 15);
        assert_eq!(s.peek(-2), 8);
    }

    #[test]
    fn test_seek() {
        let s = Store::new(|i: i32| i * i, 0_i32);
        let s2 = s.seek(6);
        assert_eq!(s2.extract(), 36);
    }

    #[test]
    fn test_extend_law() {
        // extract . extend f = f
        let s = Store::new(|i: i32| i * 2, 5_i32);
        let f = |st: &Store<i32, i32>| st.extract() + 1;
        let s2 = s.extend(f);
        assert_eq!(s2.extract(), f(&s));
    }

    #[test]
    fn test_rule30_seed() {
        // Middle cell alive, rest dead — check initial extract
        let cells = vec![false; 7];
        let mut c = cells.clone();
        c[3] = true;
        let s = make_cell_store(&c);
        assert_eq!(s.peek(3), true);
        assert_eq!(s.peek(2), false);
    }
}