[Relay] fix checkwellform (apache#2705)

* do * address comment
neo-ai · Mar 12, 2019 · e8d54a0 · e8d54a0
1 parent 27475e7
commit e8d54a0
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 5 deletions.
diff --git a/src/relay/pass/well_formed.cc b/src/relay/pass/well_formed.cc
@@ -5,14 +5,15 @@
  */
 #include <tvm/relay/pass.h>
 #include <tvm/relay/expr_functor.h>
+#include <tvm/relay/pattern_functor.h>
 #include <unordered_set>
 
 namespace tvm {
 namespace relay {
 
 
 //! brief make sure each Var is bind at most once.
-class WellFormedChecker : private ExprVisitor {
+class WellFormedChecker : private ExprVisitor, PatternVisitor {
   bool well_formed = true;
 
   std::unordered_set<Var, NodeHash, NodeEqual> s;
@@ -39,6 +40,14 @@ class WellFormedChecker : private ExprVisitor {
     CheckWellFormed(f->body);
   }
 
+  void VisitPattern(const Pattern& p) final {
+    PatternVisitor::VisitPattern(p);
+  }
+
+  void VisitVar(const Var& v) final {
+    Check(v);
+  }
+
  public:
   bool CheckWellFormed(const Expr& e) {
     this->VisitExpr(e);

diff --git a/tests/python/relay/test_ir_well_formed.py b/tests/python/relay/test_ir_well_formed.py
@@ -1,9 +1,10 @@
 import tvm
 from tvm import relay
 from tvm.relay.ir_pass import well_formed
+from tvm.relay.prelude import Prelude
 
-def test_well_formed():
-    x = relay.Var('x')
+def test_let():
+    x = relay.Var("x")
     assert well_formed(x)
     v = relay.Constant(tvm.nd.array(10))
     ty = None
@@ -18,7 +19,7 @@ def test_well_formed():
 
 
 def test_tuple():
-    x = relay.Var('x')
+    x = relay.Var("x")
     assert well_formed(x)
     v = relay.Constant(tvm.nd.array(10))
     let = relay.Let(x, v, x)
@@ -28,5 +29,23 @@ def test_tuple():
 
 
 def test_tuple_get_item():
-    t = relay.Var('t')
+    t = relay.Var("t")
     assert well_formed(relay.TupleGetItem(t, 2))
+
+
+def test_adt():
+    mod = relay.Module()
+    p = Prelude(mod)
+    x = relay.Var("x")
+    s_case = relay.Clause(relay.PatternConstructor(p.s, [relay.PatternVar(x)]), x)
+    default_case = relay.Clause(relay.PatternVar(x), x)
+    m0 = relay.Match(p.z(), [default_case])
+    m1 = relay.Match(p.z(), [s_case, default_case])
+    assert well_formed(m0)
+    assert not well_formed(m1)
+
+if __name__ == "__main__":
+    test_let()
+    test_tuple()
+    test_tuple_get_item()
+    test_adt()