Swift: statement visitor

This transfers the current state of `StmtVisitor` in the PoC, plus some
changes required for the update to swift 5.6.

Also `getLabel` in `SwiftDispatcher` got renamed to `createLabel`, and
is now correctly outputting the label assignment to the trap file.

Author: redsun82

swift/codegen/schema.yml

@@ -52,9 +52,9 @@ AstNode:
 
 ConditionElement:
   _extends: Locatable
-#  boolean: Expr?
-#  pattern: Pattern?
-#  initializer: Expr?
+  boolean: Expr?
+  pattern: Pattern?
+  initializer: Expr?
   _dir: stmt
 
 AnyFunctionType:
@@ -534,56 +534,56 @@ TypedPattern:
 
 BraceStmt:
   _extends: Stmt
-#  elements: AstNode*
+  elements: AstNode*
 
 BreakStmt:
   _extends: Stmt
-#  target_name: string?
-#  target: Stmt?
+  target_name: string?
+  target: Stmt?
 
 CaseStmt:
   _extends: Stmt
-#  body: Stmt
-#  labels: CaseLabelItem*
-#  variables: VarDecl*
+  body: Stmt
+  labels: CaseLabelItem*
+  variables: VarDecl*
 
 CaseLabelItem:
   _extends: AstNode
-#  pattern: Pattern
-#  guard: Expr?
+  pattern: Pattern
+  guard: Expr?
   _dir: stmt
 
 ContinueStmt:
   _extends: Stmt
-#  target_name: string?
-#  target: Stmt?
+  target_name: string?
+  target: Stmt?
 
 DeferStmt:
   _extends: Stmt
-#  body: BraceStmt
+  body: BraceStmt
 
 FailStmt:
   _extends: Stmt
 
 FallthroughStmt:
   _extends: Stmt
-#  fallthrough_source: CaseStmt
-#  fallthrough_dest: CaseStmt
+  fallthrough_source: CaseStmt
+  fallthrough_dest: CaseStmt
 
 LabeledStmt:
   _extends: Stmt
-#  label: string?
+  label: string?
 
 PoundAssertStmt:
   _extends: Stmt
 
 ReturnStmt:
   _extends: Stmt
-#  result: Expr?
+  result: Expr?
 
 ThrowStmt:
   _extends: Stmt
-#  sub_expr: Expr
+  sub_expr: Expr
 
 YieldStmt:
   _extends: Stmt
@@ -843,37 +843,38 @@ OverloadedDeclRefExpr:
 
 DoCatchStmt:
   _extends: LabeledStmt
-#  body: Stmt
-#  catches: CaseStmt*
+  body: Stmt
+  catches: CaseStmt*
 
 DoStmt:
   _extends: LabeledStmt
-#  body: BraceStmt
+  body: BraceStmt
 
 ForEachStmt:
   _extends: LabeledStmt
-#  body: BraceStmt
-#  sequence: Expr
-#  where: Expr?
+  pattern: Pattern
+  sequence: Expr
+  where: Expr?
+  body: BraceStmt
 
 LabeledConditionalStmt:
   _extends: LabeledStmt
-#  condition: StmtCondition
+  condition: StmtCondition
 
 StmtCondition:
   _extends: AstNode
-#  elements: ConditionElement*
+  elements: ConditionElement*
   _dir: stmt
 
 RepeatWhileStmt:
   _extends: LabeledStmt
-#  condition: Expr
-#  body: Stmt
+  condition: Expr
+  body: Stmt
 
 SwitchStmt:
   _extends: LabeledStmt
-#  expr: Expr
-#  cases: CaseStmt*
+  expr: Expr
+  cases: CaseStmt*
 
 BoundGenericClassType:
   _extends: BoundGenericType
@@ -979,16 +980,16 @@ DynamicSubscriptExpr:
 
 GuardStmt:
   _extends: LabeledConditionalStmt
-#  body: BraceStmt
+  body: BraceStmt
 
 IfStmt:
   _extends: LabeledConditionalStmt
-#  then: Stmt
-#  else: Stmt?
+  then: Stmt
+  else: Stmt?
 
 WhileStmt:
   _extends: LabeledConditionalStmt
-#  body: Stmt
+  body: Stmt
 
 AccessorDecl:
   _extends: FuncDecl

swift/extractor/SwiftDispatcher.h

@@ -67,60 +67,104 @@ class SwiftDispatcher {
   // `swift::TypeBase*`)
   TrapLabel<TypeTag> fetchLabel(swift::Type t) { return fetchLabel(t.getPointer()); }
 
+  TrapLabel<AstNodeTag> fetchLabel(swift::ASTNode node) {
+    return fetchLabelFromUnion<AstNodeTag>(node);
+  }
+
   // Due to the lazy emission approach, we must assign a label to a corresponding AST node before
   // it actually gets emitted to handle recursive cases such as recursive calls, or recursive type
   // declarations
   template <typename E>
   TrapLabelOf<E> assignNewLabel(E* e) {
     assert(waitingForNewLabel == Store::Handle{e} && "assignNewLabel called on wrong entity");
-    auto label = getLabel<TrapTagOf<E>>();
-    trap.assignStar(label);
+    auto label = createLabel<TrapTagOf<E>>();
     store.insert(e, label);
     waitingForNewLabel = std::monostate{};
     return label;
   }
 
   template <typename Tag>
-  TrapLabel<Tag> getLabel() {
-    return arena.allocateLabel<Tag>();
+  TrapLabel<Tag> createLabel() {
+    auto ret = arena.allocateLabel<Tag>();
+    trap.assignStar(ret);
+    return ret;
+  }
+
+  template <typename Tag, typename... Args>
+  TrapLabel<Tag> createLabel(Args&&... args) {
+    auto ret = arena.allocateLabel<Tag>();
+    trap.assignKey(ret, std::forward<Args>(args)...);
+    return ret;
   }
 
   template <typename Locatable>
   void attachLocation(Locatable locatable, TrapLabel<LocatableTag> locatableLabel) {
     attachLocation(&locatable, locatableLabel);
   }
 
-  // Emits a Location TRAP entry and attaches it to an AST node
+  // Emits a Location TRAP entry and attaches it to a `Locatable` trap label
   template <typename Locatable>
   void attachLocation(Locatable* locatable, TrapLabel<LocatableTag> locatableLabel) {
-    auto start = locatable->getStartLoc();
-    auto end = locatable->getEndLoc();
+    attachLocation(locatable->getStartLoc(), locatable->getEndLoc(), locatableLabel);
+  }
+
+  // Emits a Location TRAP entry for a list of swift entities and attaches it to a `Locatable` trap
+  // label
+  template <typename Locatable>
+  void attachLocation(llvm::MutableArrayRef<Locatable>* locatables,
+                      TrapLabel<LocatableTag> locatableLabel) {
+    if (locatables->empty()) {
+      return;
+    }
+    attachLocation(locatables->front().getStartLoc(), locatables->back().getEndLoc(),
+                   locatableLabel);
+  }
+
+ private:
+  // types to be supported by assignNewLabel/fetchLabel need to be listed here
+  using Store = TrapLabelStore<swift::Decl,
+                               swift::Stmt,
+                               swift::StmtCondition,
+                               swift::CaseLabelItem,
+                               swift::Expr,
+                               swift::Pattern,
+                               swift::TypeRepr,
+                               swift::TypeBase>;
+
+  void attachLocation(swift::SourceLoc start,
+                      swift::SourceLoc end,
+                      TrapLabel<LocatableTag> locatableLabel) {
     if (!start.isValid() || !end.isValid()) {
       // invalid locations seem to come from entities synthesized by the compiler
       return;
     }
     std::string filepath = getFilepath(start);
-    auto fileLabel = arena.allocateLabel<FileTag>();
-    trap.assignKey(fileLabel, filepath);
+    auto fileLabel = createLabel<FileTag>(filepath);
     // TODO: do not emit duplicate trap entries for Files
     trap.emit(FilesTrap{fileLabel, filepath});
     auto [startLine, startColumn] = sourceManager.getLineAndColumnInBuffer(start);
     auto [endLine, endColumn] = sourceManager.getLineAndColumnInBuffer(end);
-    auto locLabel = arena.allocateLabel<LocationTag>();
-    trap.assignKey(locLabel, '{', fileLabel, "}:", startLine, ':', startColumn, ':', endLine, ':',
-                   endColumn);
+    auto locLabel = createLabel<LocationTag>('{', fileLabel, "}:", startLine, ':', startColumn, ':',
+                                             endLine, ':', endColumn);
     trap.emit(LocationsTrap{locLabel, fileLabel, startLine, startColumn, endLine, endColumn});
     trap.emit(LocatablesTrap{locatableLabel, locLabel});
   }
 
- private:
-  // types to be supported by assignNewLabel/fetchLabel need to be listed here
-  using Store = TrapLabelStore<swift::Decl,
-                               swift::Stmt,
-                               swift::Expr,
-                               swift::Pattern,
-                               swift::TypeRepr,
-                               swift::TypeBase>;
+  template <typename Tag, typename... Ts>
+  TrapLabel<Tag> fetchLabelFromUnion(const llvm::PointerUnion<Ts...> u) {
+    TrapLabel<Tag> ret{};
+    assert((... || fetchLabelFromUnionCase<Tag, Ts>(u, ret)) && "llvm::PointerUnion not set");
+    return ret;
+  }
+
+  template <typename Tag, typename T, typename... Ts>
+  bool fetchLabelFromUnionCase(const llvm::PointerUnion<Ts...> u, TrapLabel<Tag>& output) {
+    if (auto e = u.template dyn_cast<T>()) {
+      output = fetchLabel(e);
+      return true;
+    }
+    return false;
+  }
 
   std::string getFilepath(swift::SourceLoc loc) {
     // TODO: this needs more testing
@@ -139,6 +183,8 @@ class SwiftDispatcher {
   // which are to be introduced in follow-up PRs
   virtual void visit(swift::Decl* decl) = 0;
   virtual void visit(swift::Stmt* stmt) = 0;
+  virtual void visit(swift::StmtCondition* cond) = 0;
+  virtual void visit(swift::CaseLabelItem* item) = 0;
   virtual void visit(swift::Expr* expr) = 0;
   virtual void visit(swift::Pattern* pattern) = 0;
   virtual void visit(swift::TypeRepr* type) = 0;

swift/extractor/SwiftTagTraits.h

@@ -35,6 +35,8 @@ using SILBoxTypeReprTag = SilBoxTypeReprTag;
   static_assert(std::is_base_of_v<TYPE##Tag, TAG>, "override is not a subtag");
 
 MAP_TAG(Stmt);
+MAP_TAG(StmtCondition);
+MAP_TAG(CaseLabelItem);
 #define ABSTRACT_STMT(CLASS, PARENT) MAP_SUBTAG(CLASS##Stmt, PARENT)
 #define STMT(CLASS, PARENT) ABSTRACT_STMT(CLASS, PARENT)
 #include "swift/AST/StmtNodes.def"

swift/extractor/SwiftVisitor.h

@@ -22,6 +22,8 @@ class SwiftVisitor : private SwiftDispatcher {
  private:
   void visit(swift::Decl* decl) override { declVisitor.visit(decl); }
   void visit(swift::Stmt* stmt) override { stmtVisitor.visit(stmt); }
+  void visit(swift::StmtCondition* cond) override { stmtVisitor.visitStmtCondition(cond); }
+  void visit(swift::CaseLabelItem* item) override { stmtVisitor.visitCaseLabelItem(item); }
   void visit(swift::Expr* expr) override { exprVisitor.visit(expr); }
   void visit(swift::Pattern* pattern) override { patternVisitor.visit(pattern); }
   void visit(swift::TypeRepr* type) override { typeReprVisitor.visit(type); }