Implement where clause variables and exists

This adds two things to where clauses:

Exists clauses are looked at in range inference but are not actually
evaluated. This allows you make things that are similarly sized to
another tensor (and will work well with the pad() operator once we
finish it):

  a(i) = 0
  where b(i) exists

Let clauses allow you to declare temporary variables whose scope
is only for the local comprehensions:

  a(i) = foo + foo*foo
  where foo = b(i)

This commit also includes a simplification of TreeView objects that
I back ported from pytorch that makes adding new tree views less
verbose.

Author: zdevito

include/tc/lang/lexer.h

@@ -83,7 +83,9 @@ namespace lang {
   _(TK_EQ, "eq", "==")                           \
   _(TK_NE, "neq", "!=")                          \
   _(TK_AND, "and", "&&")                         \
-  _(TK_OR, "or", "||")
+  _(TK_OR, "or", "||")                           \
+  _(TK_LET, "let", "")                           \
+  _(TK_EXISTS, "exists", "exists")
 
 static const char* valid_single_char_tokens = "+-*/()[]?:,={}><!";
 
@@ -380,9 +382,20 @@ struct Lexer {
     next();
     return true;
   }
+  Token lookahead() {
+    if (!lookahead_) {
+      lookahead_.reset(new Token(lex()));
+    }
+    return *lookahead_;
+  }
   Token next() {
     auto r = cur_;
-    cur_ = lex();
+    if (lookahead_) {
+      cur_ = *lookahead_;
+      lookahead_.reset();
+    } else {
+      cur_ = lex();
+    }
     return r;
   }
   void reportError(const std::string& what, const Token& t);
@@ -416,6 +429,7 @@ struct Lexer {
   }
   size_t pos;
   Token cur_;
+  std::unique_ptr<Token> lookahead_;
   SharedParserData& shared;
 };
 } // namespace lang

include/tc/lang/parser.h

@@ -149,6 +149,24 @@ struct Parser {
     auto r = parseExp();
     return RangeConstraint::create(id->range(), id, l, r);
   }
+  TreeRef parseLetBinding() {
+    auto ident = parseIdent();
+    L.expect('=');
+    auto exp = parseExp();
+    return Let::create(ident->range(), ident, exp);
+  }
+  TreeRef parseWhereClause() {
+    auto lookahead = L.lookahead();
+    if (lookahead.kind == '=') {
+      return parseLetBinding();
+    } else if (lookahead.kind == TK_IN) {
+      return parseRangeConstraint();
+    } else {
+      auto exp = parseExp();
+      L.expect(TK_EXISTS);
+      return Exists::create(exp->range(), {exp});
+    }
+  }
   TreeRef parseParam() {
     if (L.cur().kind == TK_IDENT) {
       auto ident = parseIdent();
@@ -159,10 +177,9 @@ struct Parser {
     auto ident = parseIdent();
     return Param::create(typ->range(), ident, typ);
   }
-  TreeRef parseRangeConstraints() {
+  TreeRef parseWhereClauses() {
     if (L.nextIf(TK_WHERE)) {
-      return parseNonEmptyList(
-          ',', [&](int i) { return parseRangeConstraint(); });
+      return parseNonEmptyList(',', [&](int i) { return parseWhereClause(); });
     }
     return List::create(L.cur().range, {});
   }
@@ -200,7 +217,7 @@ struct Parser {
     auto assign = parseAssignment();
     auto rhs = parseExp();
     TreeRef equivalent_statement = parseEquivalent();
-    TreeRef range_statements = parseRangeConstraints();
+    TreeRef range_statements = parseWhereClauses();
     TreeRef empty_reduction_variables = c(TK_LIST, ident->range(), {});
     return Comprehension::create(
         ident->range(),

include/tc/lang/sema.h

@@ -204,6 +204,17 @@ struct Sema {
     expectBool(exp, typeOfExpr(exp)->kind());
     return exp;
   }
+  TreeRef lookupVarOrCreateIndex(Ident ident) {
+    TreeRef type = lookup(ident, false);
+    if (!type) {
+      // variable exp is not defined, so a reduction variable is created
+      // a reduction variable index i
+      type = indexType(ident);
+      insert(index_env, ident, type, true);
+      reduction_variables.push_back(ident);
+    }
+    return type;
+  }
   TreeRef checkExp(TreeRef exp, bool allow_access) {
     switch (exp->kind()) {
       case TK_APPLY: {
@@ -250,14 +261,7 @@ struct Sema {
       } break;
       case TK_IDENT: {
         auto ident = Ident(exp);
-        TreeRef type = lookup(ident, false);
-        if (!type) {
-          // variable exp is not defined, so a reduction variable is created
-          // a reduction variable index i
-          type = indexType(exp);
-          insert(index_env, ident, type, true);
-          reduction_variables.push_back(exp);
-        }
+        auto type = lookupVarOrCreateIndex(ident);
         if (type->kind() == TK_TENSOR_TYPE) {
           auto tt = TensorType(type);
           if (tt.dims().size() != 0) {
@@ -397,6 +401,33 @@ struct Sema {
     }
     return List::create(list->range(), std::move(r));
   }
+  TreeRef checkRangeConstraint(RangeConstraint rc) {
+    // RCs are checked _before_ the rhs of the TC, so
+    // it is possible the index is not in the environment yet
+    // calling lookupOrCreate ensures it exists
+    lookupVarOrCreateIndex(rc.ident());
+    // calling looking directly in the index_env ensures that
+    // we are actually constraining an index and not some other variable
+    lookup(index_env, rc.ident(), true);
+    auto s = expectIntegral(checkExp(rc.start(), false));
+    auto e = expectIntegral(checkExp(rc.end(), false));
+    return RangeConstraint::create(rc.range(), rc.ident(), s, e);
+  }
+  TreeRef checkLet(Let l) {
+    auto rhs = checkExp(l.rhs(), true);
+    insert(let_env, l.name(), typeOfExpr(rhs), true);
+    return Let::create(l.range(), l.name(), rhs);
+  }
+  TreeRef checkWhereClause(TreeRef ref) {
+    if (ref->kind() == TK_LET) {
+      return checkLet(Let(ref));
+    } else if (ref->kind() == TK_EXISTS) {
+      auto exp = checkExp(Exists(ref).exp(), true);
+      return Exists::create(ref->range(), exp);
+    } else {
+      return checkRangeConstraint(RangeConstraint(ref));
+    }
+  }
   TreeRef checkStmt(TreeRef stmt_) {
     auto stmt = Comprehension(stmt_);
 
@@ -417,6 +448,11 @@ struct Sema {
       output_indices.push_back(new_var);
     }
 
+    // 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(
+        [&](const TreeRef& rc) { return checkWhereClause(rc); });
+
     TreeRef rhs_ = checkExp(stmt.rhs(), true);
     TreeRef scalar_type = typeOfExpr(rhs_);
 
@@ -451,14 +487,6 @@ struct Sema {
     // if we redefined an input, it is no longer valid for range expressions
     live_input_names.erase(stmt.ident().name());
 
-    auto range_constraints =
-        stmt.rangeConstraints().map([&](const RangeConstraint& rc) {
-          lookup(index_env, rc.ident(), true);
-          auto s = expectIntegral(checkExp(rc.start(), false));
-          auto e = expectIntegral(checkExp(rc.end(), false));
-          return RangeConstraint::create(rc.range(), rc.ident(), s, e);
-        });
-
     auto equivalent_statement_ =
         stmt.equivalent().map([&](const Equivalent& eq) {
           auto indices_ = eq.accesses().map(
@@ -489,10 +517,13 @@ struct Sema {
         stmt.indices(),
         stmt.assignment(),
         rhs_,
-        range_constraints,
+        where_clauses_,
         equivalent_statement_,
         reduction_variable_list);
+    // clear the per-statement environments to get ready for the next statement
     index_env.clear();
+    let_env.clear();
+
     return result;
   }
   bool isNotInplace(const TreeRef& assignment) {
@@ -538,6 +569,8 @@ struct Sema {
   }
   TreeRef lookup(const Ident& ident, bool required) {
     TreeRef v = lookup(index_env, ident, false);
+    if (!v)
+      v = lookup(let_env, ident, false);
     if (!v)
       v = lookup(env, ident, required);
     return v;
@@ -560,6 +593,7 @@ struct Sema {
 
   std::vector<TreeRef> reduction_variables; // per-statement
   Env index_env; // per-statement
+  Env let_env; // per-statement, used for where i = <exp>
 
   Env env; // name -> type
   Env annotated_output_types; // name -> type, for all annotated returns types

include/tc/lang/tree.h

@@ -75,34 +75,19 @@ struct Tree : std::enable_shared_from_this<Tree> {
   virtual TreeRef map(std::function<TreeRef(TreeRef)> fn) {
     return shared_from_this();
   }
-  template <typename... Args>
-  void match(int k, Args&... args) {
-    matchD(k, "unknown", 0, args...);
+  void expect(int k) {
+    expect(k, trees().size());
   }
-  template <typename... Args>
-  void matchD(int k, const char* filename, int lineno, Args&... args) {
-    if (kind() != k) {
+  void expect(int k, int numsubtrees) {
+    if (kind() != k || trees().size() != numsubtrees) {
       std::stringstream ss;
-      ss << filename << ":" << lineno << ": expecting kind '" << kindToString(k)
-         << "' but found '" << kind() << "'\n";
+      ss << "expected kind '" << kindToString(k) << "' with " << numsubtrees
+         << " subtrees but found '" << kindToString(kind()) << "' with "
+         << trees().size() << " subtrees.\n";
       range().highlight(ss);
       throw std::runtime_error(ss.str());
     }
-    std::initializer_list<TreeRef*> vars = {&args...};
-    if (vars.size() > trees().size()) {
-      std::stringstream ss;
-      ss << filename << ":" << lineno << ": trying to match " << vars.size()
-         << " variables against " << trees().size() << " values in list.\n";
-      range().highlight(ss);
-      throw std::runtime_error(ss.str());
-    }
-    size_t i = 0;
-    for (TreeRef* v : vars) {
-      *v = trees()[i++];
-    }
   }
-
- private:
   int kind_;
 };