teach sema to extract iteration variables from LHS subtrees

Sema needs the list of iteration variables on the Comprehension LHS to
differentiate between reduction and non-reduction variables, the former
appearing only on the RHS.  Original implementation assumes
Comprehension LHS is a tensor whose indices are Idents and ignores more
complex constructs.  With indirection support, comprehensions like
O(A(i)) = B(i) are possible but i is interpreted as a reduction
dimension by Sema.

Traverse indices of the LHS Tensor in Comprehension recursively,
inspecting subtrees of Access and Apply trees and collecting all Idents.

Author: ftynse

tc/lang/sema.h

@@ -437,17 +437,35 @@ struct Sema {
       return checkRangeConstraint(RangeConstraint(ref));
     }
   }
+
+ private:
+  // Traverse the list of trees, recursively descending into arguments of APPLY
+  // and ACCESS subtrees and into all subtrees of different types (mostly
+  // expressions), and collect names and types of IDENT subtrees in
+  // "index_env".  Expects to be called on the indices of the LHS tensor.
+  template <typename Collection>
+  void registerLHSIndices(const Collection& treeRefs) {
+    for (const auto& treeRef : treeRefs) {
+      if (treeRef->kind() == TK_IDENT) {
+        std::string idx = Ident(treeRef).name();
+        auto typ = indexType(treeRef);
+        insert(index_env, Ident(treeRef), typ, true);
+      } else if (treeRef->kind() == TK_APPLY) {
+        registerLHSIndices(Apply(treeRef).arguments());
+      } else if (treeRef->kind() == TK_ACCESS) {
+        registerLHSIndices(Access(treeRef).arguments());
+      } else {
+        registerLHSIndices(treeRef->trees());
+      }
+    }
+  }
+
+ public:
   TreeRef checkStmt(TreeRef stmt_) {
     auto stmt = Comprehension(stmt_);
 
     // register index variables (non-reductions)
-    for (const auto& index : stmt.indices()) {
-      if (index->kind() == TK_IDENT) {
-        std::string idx = Ident(index).name();
-        auto typ = indexType(index);
-        insert(index_env, Ident(index), typ, true);
-      }
-    }
+    registerLHSIndices(stmt.indices());
 
     // make dimension variables for each dimension of the output tensor
     std::string name = stmt.ident().name();
@@ -464,6 +482,7 @@ struct Sema {
 
     // where clauses are checked _before_ the rhs because they
     // introduce let bindings that are in scope for the rhs
+    //
     auto where_clauses_ = stmt.whereClauses().map(
         [&](TreeRef rc) { return checkWhereClause(rc); });
 

test/test_core.cc

@@ -346,6 +346,19 @@ def gather(int32(N) A, int32(N) B) -> (O) {
   CHECK_EQ(scop->mustWrites, scop->mayWrites);
 }
 
+TEST_F(TC2Isl, Computed) {
+  string tc = R"TC(
+def gather(int32(N) A, int32(N) B) -> (O) {
+    O(i - 2) = A(i) + B(i)
+}
+)TC";
+  auto scop = MakeScop(tc);
+  CHECK(!scop->mustWrites.is_empty())
+      << "expected non-empty must-writes for gather, got\n"
+      << scop->mustWrites;
+  CHECK_EQ(scop->mustWrites, scop->mayWrites);
+}
+
 // FIXME: range inference seems unaware of indirections on the LHS
 TEST_F(TC2Isl, DISABLED_MustWritesSubsetMayWrites) {
   string tc = R"TC(