Fleshed out NN context interface

3 files changed, 108 insertions, 11 deletions

diff --git a/src/ctx.rs b/src/ctx.rs
new file mode 100644
index 0000000..38aa4ae
--- /dev/null
+++ b/src/ctx.rs
@@ -0,0 +1,90 @@
+use crate::quad::{Quad, Node};
+
+/// Represents a context with shared state.
+pub struct Ctx();
+
+impl Ctx {
+    /// Creates a new uninitialized context.
+    pub fn new_empty() -> Self {
+        Ctx()
+    }
+
+    /// Combines 4 child node representations into a single representation
+    /// Using a neural network.
+    pub fn combine<B>(&mut self, compr: [B; 4]) -> B {
+        todo!("Build new B from 4 child B");
+    }
+
+    /// Compresses a base-level cell into a vector.
+    pub fn compress_base<A, B>(&mut self, base: A) -> B {
+        todo!("Turn Base Cell into a vector B");
+    }
+
+    /// Compresses a single node into a vector representation.
+    /// Returns `None` if node has already been compressed and trimmed from tree.
+    /// To recover a trimmed node, use `expand` on the compressed representation.
+    pub fn compress<A: Default, B: Copy>(&mut self, quad: &Quad<A, B>) -> Option<B> {
+        match quad {
+            Quad::Base(b) => Some(self.compress_base(b)),
+            Quad::Node(n) => Some(
+                self.combine([
+                    n[0].compr,
+                    n[1].compr,
+                    n[2].compr,
+                    n[3].compr,
+                ])
+            ),
+            Quad::Cached => None,
+        }
+    }
+
+    /// Compresses a base-level cell into a vector.
+    fn expand_base<A: Default, B: Copy>(&mut self, compr: B) -> A {
+        todo!("Turn compressed B into the A that made it");
+    }
+
+    fn expand_node<B>(&mut self, compr: B) -> [B; 4] {
+        todo!("Turn compressed B into 4 child B that made it");
+    }
+
+    /// Expands the compressed representation of a node into a node with 4 children.
+    pub fn expand<A: Default, B: Copy>(&mut self, mut compr: Node<A, B>) -> Node<A, B> {
+        match compr.data {
+            // Can't expand a base node.
+            Quad::Base(_) => {
+                debug_assert!(compr.depth == 0, "Tree is malformed at the leaves");
+                compr
+            },
+
+            // No-op if node is already expanded.
+            Quad::Node(_) => {
+                debug_assert!(compr.depth != 0, "Tree is malformed along the trunk");
+                compr
+            },
+
+            Quad::Cached => {
+                // Expand and repace the current node data.
+                compr.data = if compr.depth == 0 {
+                    // Base case.
+                    Quad::Base(self.expand_base(compr.compr))
+                } else {
+                    // Expand the children into their corresponding vectors.
+                    let c = self.expand_node(compr.compr);
+                    // Pack the children into a new Node.
+                    let new_depth = compr.depth - 1;
+                    let children = [
+                        Node::new_cached(c[0], new_depth),
+                        Node::new_cached(c[1], new_depth),
+                        Node::new_cached(c[2], new_depth),
+                        Node::new_cached(c[3], new_depth),
+                    ];
+                    // Heap-allocate the node lol.
+                    Quad::Node(Box::new(children))
+                };
+
+                // The updated node.
+                compr
+            },
+        }
+    }
+}
diff --git a/src/main.rs b/src/main.rs
index 5abe13f..1757013 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,4 +1,5 @@
 mod quad;
+mod ctx;
 
 fn main() {
     println!("Hello, world!");
diff --git a/src/quad.rs b/src/quad.rs
index 583e5a9..bab0491 100644
--- a/src/quad.rs
+++ b/src/quad.rs
@@ -1,6 +1,8 @@
+use crate::ctx::Ctx;
+
 /// Represents the data present in a quad-tree node.
 /// May be the base-level repr, or a node with 4 children.
-pub enum Quad<A: Default, B> {
+pub enum Quad<A: Default, B: Copy> {
     /// Base cell in grid.
     /// May actually be a chunk of cells for performance.
     Base(A),
@@ -14,21 +16,18 @@ pub enum Quad<A: Default, B> {
 /// Has a depth denoting the number of nodes below it.
 /// Nodes should only be siblings of nodes with the same depth.
 /// Data stored inside a quadtree node, including children, are in `data`.
-pub struct Node<A: Default, B> {
-    depth: usize,
-    compr: B,
-    data:  Quad<A, B>,
+pub struct Node<A: Default, B: Copy> {
+    pub depth: usize,
+    pub compr: B,
+    pub data:  Quad<A, B>,
 }
 
-/// Represents a context with shared state.
-pub struct Ctx();
-
-impl<A: Default, B> Node<A, B> {
+impl<A: Default, B: Copy> Node<A, B> {
     /// Creates a new tree from a single base node
     pub fn new_base(base: A, ctx: &mut Ctx) -> Self {
         Node {
             depth: 0,
-            compr: ctx.compress(&base),
+            compr: ctx.compress_base(&base),
             data:  Quad::Base(base),
         }
     }
@@ -38,11 +37,18 @@ impl<A: Default, B> Node<A, B> {
         Self::new_base(Default::default(), ctx)
     }
 
+    pub fn new_cached(compr: B, depth: usize) -> Self {
+        Node {
+            depth,
+            compr,
+            data: Quad::Cached,
+        }
+    }
+
     /// Creates a new node double the size by centering the current node
     /// on a node double the size.
     pub fn pad_empty(self, ctx: &mut Ctx) -> Self {
         todo!()
     }
 
-    
 }