some comments

Makefile

@@ -176,7 +176,7 @@ arb_os/arbos-upgrade-base.mexe: $(arbos_source_no_bridge) .make/tools
 	exit 1
 	@touch .make/libs
 
-.make/fmt: src/**.rs Cargo.* .make/install
+.make/fmt: src/*.rs src/*/*.rs Cargo.* .make/install
 	cargo fmt
 	@touch .make/fmt
 

src/compile/codegen.rs

@@ -35,11 +35,13 @@ struct Codegen<'a> {
     globals: &'a HashMap<StringId, GlobalVar>,
     /// Whether to elide debug-only constructs like assert().
     release_build: bool,
-
+    /// The open set of scopes
     scopes: Vec<Scope>,
-    next_slot: SlotNum,
+    /// The next slot available for assignment
+    next_assignable_slot: SlotNum,
 }
 
+/// Represents a mini scope and the values it has access to.
 #[derive(Clone, Debug, Default)]
 struct Scope {
     /// A variable's current slot
@@ -49,13 +51,14 @@ struct Scope {
 }
 
 impl Codegen<'_> {
+    /// Get the next available slot.
     fn next_slot(&mut self) -> SlotNum {
-        let next = self.next_slot;
-        self.next_slot += 1;
+        let next = self.next_assignable_slot;
+        self.next_assignable_slot += 1;
         next
     }
 
-    /// Create a new scope, inheriting the locals of the one prior
+    /// Create a new scope, inheriting the locals of the one prior.
     fn open_scope(&mut self) {
         let mut scope = match self.scopes.last() {
             Some(scope) => scope.clone(),
@@ -65,30 +68,33 @@ impl Codegen<'_> {
         self.scopes.push(scope);
     }
 
+    /// Create a new assignment for a local variable.
     fn set_local(&mut self, local: StringId, slot: SlotNum) {
-        let len = self.scopes.len();
-        self.scopes[len - 1].locals.insert(local, slot);
+        let last = self.scopes.last_mut().expect("no scope");
+        last.locals.insert(local, slot);
     }
 
+    /// Shadow a variable, saving the last unshadowed assignment for phi-ing if needed.
     fn shadow(&mut self, local: StringId, slot: SlotNum) {
-        let len = self.scopes.len();
+        let last = self.scopes.last_mut().expect("no scope");
 
         // for the first time shadowing, we save the old value if it exists
-        if let Some(old) = self.scopes[len - 1].locals.get(&local) {
+        if let Some(old) = last.locals.get(&local) {
             let old = *old;
-            if !self.scopes[len - 1].shadows.contains_key(&local) {
-                self.scopes[len - 1].shadows.insert(local, old);
+            if !last.shadows.contains_key(&local) {
+                last.shadows.insert(local, old);
             }
         }
-
-        self.scopes[len - 1].locals.insert(local, slot);
+        last.locals.insert(local, slot);
     }
 
+    /// Get the currently accessible slot assignment for a variable in scope.
     fn get_local(&mut self, local: &StringId) -> Option<SlotNum> {
-        let len = self.scopes.len();
-        self.scopes[len - 1].locals.get(local).cloned()
+        let last = self.scopes.last_mut().expect("no scope");
+        last.locals.get(local).cloned()
     }
 
+    /// Debug print the open scope's current assignments.
     fn _print_locals(&self, title: &str, depth: usize) {
         let len = self.scopes.len();
         let scope = &self.scopes[len - 1];
@@ -128,7 +134,8 @@ impl Codegen<'_> {
 /// and globals lists the globals available to the func.
 ///
 /// Each func gets a unique, hashed label id, after which local labels are assigned. This ensures
-/// two labels are the same iff they point to the same destination.
+/// two labels are the same iff they point to the same destination. func_labels maps local `StringId`s
+/// to these globally consistent labels.
 pub fn mavm_codegen_func(
     mut func: TypeCheckedFunc, // only mutable because of child_nodes()
     string_table: &StringTable,
@@ -200,7 +207,7 @@ pub fn mavm_codegen_func(
         globals,
         release_build,
         scopes: vec![Scope::default()],
-        next_slot: 0,
+        next_assignable_slot: 0,
     };
 
     let mut declare = vec![];
@@ -209,13 +216,18 @@ pub fn mavm_codegen_func(
 
     codegen(func.child_nodes(), &mut cgen, 0, declare)?;
 
-    let space_for_locals = cgen.next_slot;
+    let space_for_locals = cgen.next_assignable_slot;
 
     code[make_frame_offset] = opcode!(@MakeFrame(space_for_locals, prebuilt));
 
     Ok((code, label_gen, space_for_locals))
 }
 
+/// Codegen a scope of typechecked nodes.
+///
+/// stack_items counts the number of items that need be popped for an early return.
+/// declare carries a list of variables to immediately shadow upon opening the scope,
+/// which is useful for function arguments and if-let.
 fn codegen(
     nodes: Vec<TypeCheckedNode>,
     cgen: &mut Codegen,
@@ -431,7 +443,7 @@ fn codegen(
                         // if-let is tricky since the local variable isn't defined in the same scope.
                         // To work around this, we get the next slot without advancing. This means
                         // not actually *calling* next_slot().
-                        let slot = cgen.next_slot;
+                        let slot = cgen.next_assignable_slot;
                         cgen.code.push(opcode!(@SetLocal(slot)));
                         block!(block, vec![*id]);
                         cgen.code.push(opcode!(Jump, Value::Label(end_label)));
@@ -821,7 +833,6 @@ fn codegen(
     }
 
     let debug = cgen.code.last().unwrap().debug_info;
-
     let scope = cgen.scopes.pop().expect("No scope");
 
     for (local, mut slot) in scope.locals {

src/compile/translate.rs

@@ -137,6 +137,7 @@ pub fn read_capture_data(code: Vec<Instruction>) -> (Vec<Instruction>, ClosureAs
     (out, captures)
 }
 
+/// Use capture annotations to bridge the gap between packing a closure and calling it.
 pub fn pack_closures(
     code: &Vec<Instruction>,
     capture_map: &HashMap<LabelId, ClosureAssignments>,

src/optimize/mod.rs

@@ -14,6 +14,7 @@ use rand::prelude::*;
 use rand::rngs::SmallRng;
 use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet, VecDeque};
 
+/// Represents a block of instructions that has no control flow
 pub enum BasicBlock {
     Code(Vec<Instruction>),
     Meta(&'static str),
@@ -34,6 +35,7 @@ impl BasicBlock {
     }
 }
 
+/// Represents control flow between two blocks.
 pub enum BasicEdge {
     Forward,
     Jump,
@@ -68,7 +70,7 @@ impl BasicGraph {
 
         let should_print = code.iter().any(|x| x.debug_info.attributes.codegen_print);
 
-        // Create basic blocks by splitting on jumps and labels
+        // create basic blocks by splitting on jumps and labels
         let mut block_data = vec![];
         for curr in code {
             match curr.opcode {
@@ -119,7 +121,7 @@ impl BasicGraph {
             }
         }
 
-        // Add jump edges
+        // add jump edges
         for node in nodes {
             let last = match &graph[node] {
                 BasicBlock::Code(code) => match code.iter().last().cloned() {
@@ -313,11 +315,25 @@ impl BasicGraph {
 
     /// Efficiently assign frame slots to minimize the frame size.
     pub fn color(&mut self, frame_size: FrameSize) {
+        // Algorithm
+        //   Determine which values would overwrite each other if reassigned the same slot.
+        //   Reassign values using as few slots as possible given this knowledge
+        //
+        // Steps
+        //   - Create a graph of which values are phi'd together.
+        //   - Build a conflict graph whose nodes represent values in the local function frame
+        //     and whose edges represent the inability to give each the same slot.
+        //   - Contract phi nodes so that each side gets the same slot assignment.
+        //   - Color the contracted graph with as few colors as best we can.
+        //   - Color the original graph 1-to-1 using the contracted graph.
+        //   - Modify each instruction to reflect the new slot assignments.
+
         let mut conflicts: StableGraph<SlotNum, (), Undirected> = StableGraph::default();
         let mut phi_graph: UnGraph<SlotNum, ()> = UnGraph::default();
 
         for slot in 0..frame_size {
-            // we do this so that slot & index are the same
+            // we do this so that slot & index are interchangeable.
+            // conflicts is a stable graph to preserve this property when we remove nodes.
             conflicts.add_node(slot);
             phi_graph.add_node(slot);
         }
@@ -330,7 +346,7 @@ impl BasicGraph {
             }
         }
 
-        // Walking the graph backwards is heuristically faster.
+        // Walking the graph backwards is heuristically faster by order O(n).
         let mut work_list: VecDeque<_> = self.graph.node_indices().rev().collect();
         let mut work_set: HashSet<_> = work_list.clone().into_iter().collect();
         let mut node_needs = HashMap::new();
@@ -425,7 +441,11 @@ impl BasicGraph {
             }
         }
 
-        // Coloring a graph is NP-complete
+        // Coloring a graph its chromatic number is NP-complete.
+        // We find the following algorithm approximates this well for the kinds of graphs
+        // SSA tends to produce. How this works is we successively pick the least-used color
+        // we can for each node in the graph. When no color is available, we make a new one.
+        // We repeat this process against a random ording some number of times and take the best.
         let mut rng = SmallRng::seed_from_u64(0);
         let mut best_assignments = HashMap::new();
         let mut ncolors = usize::MAX;
@@ -464,7 +484,8 @@ impl BasicGraph {
             }
         }
 
-        // apply colors to phi'd slots
+        // Apply colors to phi'd slots. Since each came from a strongly connected
+        // component, no two Dfs's will re-assign the same node.
         for node in conflicts.node_indices() {
             let slot = node.index() as u32;
             let color = *best_assignments.get(&slot).unwrap();
@@ -494,6 +515,7 @@ impl BasicGraph {
             }
         }
 
+        // The frame should be much smaller now, so it makes sense to shrink it.
         self.shrink_frame();
 
         if self.should_print {