From 57476b110274c237081eb285bec72c46fe0731e0 Mon Sep 17 00:00:00 2001
From: Alex Reinking <areinking@adobe.com>
Date: Fri, 13 Dec 2024 03:36:02 -0500
Subject: [PATCH] Add more detailed documentation to add_let.

---
 src/Func.cpp | 42 ++++++++++++++++++++++++++++++++----------
 1 file changed, 32 insertions(+), 10 deletions(-)

diff --git a/src/Func.cpp b/src/Func.cpp
index 7ea846490ed3..b6fc534b3dbc 100644
--- a/src/Func.cpp
+++ b/src/Func.cpp
@@ -635,18 +635,40 @@ optional<Dim> find_dim(const vector<Dim> &items, const VarOrRVar &v) {
 
 using SubstitutionMap = std::map<string, Expr>;
 
-// This is a helper function for building up a substitution map that
-// corresponds to pushing down a nest of lets. The lets should be fed
-// to this function from innermost to outermost. This is equivalent to
-// building a let-nest as a one-hole context and then simplifying.
-void add_let(SubstitutionMap &proj, const string &name, const Expr &value) {
-    internal_assert(!proj.count(name)) << "would shadow " << name << " in let nest.\n"
-                                       << "\tPresent value: " << proj[name] << "\n"
-                                       << "\tProposed value: " << value;
-    for (auto &[_, e] : proj) {
+/** This is a helper function for building up a substitution map that
+ * corresponds to pushing down a nest of lets. The lets should be fed
+ * to this function from innermost to outermost. This is equivalent to
+ * building a let-nest as a one-hole context and then simplifying.
+ *
+ * This looks like it might be quadratic or worse, and technically it is,
+ * but this isn't a problem for the way it is used inside rfactor. There
+ * are only a few uses:
+ *
+ *   1. Remapping preserved RVars to new RVars
+ *   2. Remapping factored RVars to new Vars
+ *   3. Filling the holes in the associative template
+ *   4. Accumulating the lets from ApplySplit
+ *
+ * These are naturally bounded by O(#splits + #dims) which is quite small
+ * in practice. Classes (1) and (2) cannot blow up expressions since they
+ * simply rename variables. Class (3) cannot blow up expressions either
+ * since nothing else can refer to the holes. That leaves only class (1).
+ * Fortunately, the lets generated by splits are benign. Split factors can't
+ * refer to RVars, and we won't see the consumed RVars in another split. So
+ * in total, this avoids any sort of exponentially sized substitution.
+ *
+ * @param subst The existing let nest (represented by a SubstitutionMap).
+ * @param name The name to bind, cannot already exist in the nest.
+ * @param value The value to bind. Will be substituted into nested values.
+ */
+void add_let(SubstitutionMap &subst, const string &name, const Expr &value) {
+    internal_assert(!subst.count(name)) << "would shadow " << name << " in let nest.\n"
+                                        << "\tPresent value: " << subst[name] << "\n"
+                                        << "\tProposed value: " << value;
+    for (auto &[_, e] : subst) {
         e = substitute(name, value, e);
     }
-    proj.emplace(name, value);
+    subst.emplace(name, value);
 }
 
 pair<ReductionDomain, SubstitutionMap> project_rdom(const vector<Dim> &dims, const ReductionDomain &rdom, const vector<Split> &splits) {