From 562af2877316b5eb4cdb246acdbb5b453b31bb70 Mon Sep 17 00:00:00 2001
From: Jim Blandy <jimb@red-bean.com>
Date: Tue, 6 Jul 2021 11:13:28 -0700
Subject: [PATCH] [spv-out] Move cached expression table to BlockContext.

Since the table of cached expressions is only meaningful within a single
function, it's really something that should only be accessed from
`BlockContext`.

However, to save heap allocations, it makes sense to retain it in the `Writer`
between functions. But the `Writer` field should have a different name, to
ensure people don't use it by accident.
---
 src/back/spv/block.rs  |  6 +++---
 src/back/spv/mod.rs    | 19 +++++++++++++++++--
 src/back/spv/writer.rs | 16 +++++++++++-----
 3 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/src/back/spv/block.rs b/src/back/spv/block.rs
index f45203076d..37d5ebfd81 100644
--- a/src/back/spv/block.rs
+++ b/src/back/spv/block.rs
@@ -197,7 +197,7 @@ impl<'w> BlockContext<'w> {
 
             // The chain rule: if this `Access...`'s `base` operand was
             // previously omitted, then omit this one, too.
-            _ => self.writer.cached.ids[expr_handle.index()] == 0,
+            _ => self.cached.ids[expr_handle.index()] == 0,
         }
     }
 
@@ -1150,7 +1150,7 @@ impl<'w> BlockContext<'w> {
             crate::Expression::ArrayLength(expr) => self.write_runtime_array_length(expr, block)?,
         };
 
-        self.writer.cached[expr_handle] = id;
+        self.cached[expr_handle] = id;
         Ok(())
     }
 
@@ -1578,7 +1578,7 @@ impl<'w> BlockContext<'w> {
 
                     let type_id = match result {
                         Some(expr) => {
-                            self.writer.cached[expr] = id;
+                            self.cached[expr] = id;
                             self.writer.lookup_function_call.insert(expr, id);
                             let ty_handle = self.ir_module.functions[local_function]
                                 .result
diff --git a/src/back/spv/mod.rs b/src/back/spv/mod.rs
index 952fb72041..ffbbf10005 100644
--- a/src/back/spv/mod.rs
+++ b/src/back/spv/mod.rs
@@ -268,6 +268,14 @@ enum Dimension {
     Matrix,
 }
 
+/// A map from evaluated [`Expression`s] to their SPIR-V ids.
+///
+/// When we emit code to evaluate a given `Expression`, we record the
+/// SPIR-V id of its value here, under its `Handle<Expression>` index.
+///
+/// A `CachedExpressions` value can be indexed by a `Handle<Expression>` value.
+///
+/// [emit]: index.html#expression-evaluation-time-and-scope
 #[derive(Default)]
 struct CachedExpressions {
     ids: Vec<Word>,
@@ -338,6 +346,9 @@ struct BlockContext<'w> {
     /// The [`back::spv::Function`] to which we are contributing SPIR-V instructions.
     function: &'w mut Function,
 
+    /// SPIR-V ids for expressions we've evaluated.
+    cached: CachedExpressions,
+
     /// The `Writer`'s temporary vector, for convenience.
     temp_list: Vec<Word>,
 }
@@ -357,7 +368,7 @@ impl BlockContext<'_> {
     }
 
     fn cached(&self, expression: Handle<crate::Expression>) -> Word {
-        self.writer.cached[expression]
+        self.cached[expression]
     }
 }
 
@@ -386,7 +397,11 @@ pub struct Writer {
     constant_ids: Vec<Word>,
     cached_constants: crate::FastHashMap<(crate::ScalarValue, crate::Bytes), Word>,
     global_variables: Vec<GlobalVariable>,
-    cached: CachedExpressions,
+
+    // Cached expressions are only meaningful within a BlockContext, but we
+    // retain the table here between functions to save heap allocations.
+    saved_cached: CachedExpressions,
+
     gl450_ext_inst_id: Word,
     // Just a temporary list of SPIR-V ids
     temp_list: Vec<Word>,
diff --git a/src/back/spv/writer.rs b/src/back/spv/writer.rs
index 89e7146590..8346e8f42e 100644
--- a/src/back/spv/writer.rs
+++ b/src/back/spv/writer.rs
@@ -82,7 +82,7 @@ impl Writer {
             constant_ids: Vec::new(),
             cached_constants: crate::FastHashMap::default(),
             global_variables: Vec::new(),
-            cached: CachedExpressions::default(),
+            saved_cached: CachedExpressions::default(),
             gl450_ext_inst_id,
             temp_list: Vec::new(),
         })
@@ -131,7 +131,7 @@ impl Writer {
             constant_ids: take(&mut self.constant_ids).recycle(),
             cached_constants: take(&mut self.cached_constants).recycle(),
             global_variables: take(&mut self.global_variables).recycle(),
-            cached: take(&mut self.cached).recycle(),
+            saved_cached: take(&mut self.saved_cached).recycle(),
             temp_list: take(&mut self.temp_list).recycle(),
         };
 
@@ -429,13 +429,16 @@ impl Writer {
             ir_function,
             fun_info: info,
             function: &mut function,
+            // Re-use the cached expression table from prior functions.
+            cached: std::mem::take(&mut self.saved_cached),
+
             // Steal the Writer's temp list for a bit.
             temp_list: std::mem::take(&mut self.temp_list),
             writer: self,
         };
 
         // fill up the pre-emitted expressions
-        context.writer.cached.reset(ir_function.expressions.len());
+        context.cached.reset(ir_function.expressions.len());
         for (handle, expr) in ir_function.expressions.iter() {
             if expr.needs_pre_emit() {
                 context.cache_expression_value(handle, &mut prelude)?;
@@ -449,8 +452,11 @@ impl Writer {
         context.write_block(main_id, &ir_function.body, None, LoopContext::default())?;
 
         // Consume the `BlockContext`, ending its borrows and letting the
-        // `Writer` steal back its temp_list.
-        let BlockContext { temp_list, .. } = context;
+        // `Writer` steal back its cached expression table and temp_list.
+        let BlockContext {
+            cached, temp_list, ..
+        } = context;
+        self.saved_cached = cached;
         self.temp_list = temp_list;
 
         function.to_words(&mut self.logical_layout.function_definitions);