1 files changed, 15 insertions, 3 deletions
diff --git a/src/passes/Directize.cpp b/src/passes/Directize.cpp
index 0bb5aebd9..34930d127 100644
--- a/src/passes/Directize.cpp
+++ b/src/passes/Directize.cpp
@@ -251,9 +251,21 @@ struct Directize : public Pass {
         if (func->imported()) {
           return;
         }
-        for (auto* set : FindAll<TableSet>(func->body).list) {
-          tablesWithSet.insert(set->table);
-        }
+
+        struct Finder : public PostWalker<Finder> {
+          TablesWithSet& tablesWithSet;
+
+          Finder(TablesWithSet& tablesWithSet) : tablesWithSet(tablesWithSet) {}
+
+          void visitTableSet(TableSet* curr) {
+            tablesWithSet.insert(curr->table);
+          }
+          void visitTableFill(TableFill* curr) {
+            tablesWithSet.insert(curr->table);
+          }
+        };
+
+        Finder(tablesWithSet).walkFunction(func);
       });
 
     for (auto& [_, names] : analysis.map) {