diff --git a/src/relay/analysis/util.cc b/src/relay/analysis/util.cc index b681b90d58e8..f07c14a286d8 100644 --- a/src/relay/analysis/util.cc +++ b/src/relay/analysis/util.cc @@ -214,9 +214,13 @@ class VarVisitor : protected ExprVisitor, protected PatternVisitor { } void VisitExpr_(const LetNode* op) final { - MarkBounded(op->var); - VisitExpr(op->value); - VisitExpr(op->body); + Expr let = GetRef(op); + while (auto let_node = let.as()) { + MarkBounded(let_node->var); + VisitExpr(let_node->value); + let = let_node->body; + } + VisitExpr(let); } void VisitPattern(const Pattern& p) final { PatternVisitor::VisitPattern(p); } diff --git a/src/relay/analysis/well_formed.cc b/src/relay/analysis/well_formed.cc index 16f8285c83b3..bccb1e1bf1ce 100644 --- a/src/relay/analysis/well_formed.cc +++ b/src/relay/analysis/well_formed.cc @@ -74,12 +74,21 @@ class WellFormedChecker : private ExprVisitor, PatternVisitor { } void VisitExpr_(const LetNode* l) final { - Scope s(this); - // we do letrec only for FunctionNode, - // but shadowing let in let binding is likely programming error, and we should forbidden it. - Bound(l->var); - CheckWellFormed(l->value); - CheckWellFormed(l->body); + std::vector scopes; + Expr let = GetRef(l); + while (auto let_node = let.as()) { + scopes.push_back(new Scope(this)); + // we do letrec only for FunctionNode, + // but shadowing let in let binding is likely programming error, and we should forbidden it. + Bound(let_node->var); + CheckWellFormed(let_node->value); + let = let_node->body; + } + CheckWellFormed(let); + while (!scopes.empty()) { + delete scopes.back(); + scopes.pop_back(); + } } void VisitExpr_(const FunctionNode* f) final {