Merge pull request #68210 from hamishknight/for-the-sake-of-argument

Author: hamishknight

include/swift/AST/ASTWalker.h

@@ -20,6 +20,7 @@
 
 namespace swift {
 
+class Argument;
 class ArgumentList;
 class Decl;
 class Expr;
@@ -643,6 +644,27 @@ class ASTWalker {
     return Action::Continue(ArgList);
   }
 
+  /// This method is called when first visiting an argument in an argument list,
+  /// before walking into its expression.
+  ///
+  /// \param Arg The argument to walk.
+  ///
+  /// \returns The walking action to perform.
+  ///
+  /// The default implementation returns \c Action::Continue().
+  virtual PreWalkAction walkToArgumentPre(const Argument &Arg) {
+    return Action::Continue();
+  }
+
+  /// This method is called after visiting an argument in an argument list.
+  ///
+  /// \returns The walking action to perform.
+  ///
+  /// The default implementation returns \c Action::Continue().
+  virtual PostWalkAction walkToArgumentPost(const Argument &Arg) {
+    return Action::Continue();
+  }
+
 protected:
   ASTWalker() = default;
   ASTWalker(const ASTWalker &) = default;

lib/AST/ASTWalker.cpp

@@ -1635,12 +1635,26 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
     return false;
   }
 
+  [[nodiscard]]
+  bool doIt(ArgumentList *ArgList, unsigned Idx) {
+    auto Arg = ArgList->get(Idx);
+    return traverse(
+        Walker.walkToArgumentPre(Arg),
+        [&]() {
+          auto *E = doIt(Arg.getExpr());
+          if (!E)
+            return true;
+          ArgList->setExpr(Idx, E);
+          return false;
+        },
+        [&]() { return Walker.walkToArgumentPost(Arg); });
+  }
+
   [[nodiscard]]
   ArgumentList *visit(ArgumentList *ArgList) {
     for (auto Idx : indices(*ArgList)) {
-      auto *E = doIt(ArgList->getExpr(Idx));
-      if (!E) return nullptr;
-      ArgList->setExpr(Idx, E);
+      if (doIt(ArgList, Idx))
+        return nullptr;
     }
     return ArgList;
   }

lib/IDE/SyntaxModel.cpp

@@ -372,13 +372,6 @@ class ModelASTWalker : public ASTWalker {
   friend class InactiveClauseRAII;
   bool inInactiveClause = false;
 
-  struct ParentArgsTy {
-    Expr *Parent = nullptr;
-    llvm::DenseMap<Expr *, Argument> Args;
-  };
-  /// A mapping of argument expressions to their full argument info.
-  SmallVector<ParentArgsTy, 4> ParentArgs;
-
 public:
   SyntaxModelWalker &Walker;
   ArrayRef<SyntaxNode> TokenNodes;
@@ -401,10 +394,7 @@ class ModelASTWalker : public ASTWalker {
 
   void visitSourceFile(SourceFile &SrcFile, ArrayRef<SyntaxNode> Tokens);
 
-  PreWalkResult<ArgumentList *>
-  walkToArgumentListPre(ArgumentList *ArgList) override;
-  PostWalkResult<ArgumentList *>
-  walkToArgumentListPost(ArgumentList *ArgList) override;
+  PreWalkAction walkToArgumentPre(const Argument &Arg) override;
 
   PreWalkResult<Expr *> walkToExprPre(Expr *E) override;
   PostWalkResult<Expr *> walkToExprPost(Expr *E) override;
@@ -544,69 +534,38 @@ static bool shouldTreatAsSingleToken(const SyntaxStructureNode &Node,
              SM.getLineAndColumnInBuffer(Node.Range.getEnd()).first;
 }
 
-ASTWalker::PreWalkResult<ArgumentList *>
-ModelASTWalker::walkToArgumentListPre(ArgumentList *ArgList) {
-  Expr *ParentExpr = Parent.getAsExpr();
-  if (!ParentExpr)
-    return Action::Continue(ArgList);
-
-  ParentArgsTy Mapping;
-  Mapping.Parent = ParentExpr;
-  for (auto Arg : *ArgList) {
-    auto res = Mapping.Args.try_emplace(Arg.getExpr(), Arg);
-    assert(res.second && "Duplicate arguments?");
-    (void)res;
-  }
-  ParentArgs.push_back(std::move(Mapping));
-  return Action::Continue(ArgList);
-}
+ASTWalker::PreWalkAction
+ModelASTWalker::walkToArgumentPre(const Argument &Arg) {
+  if (isVisitedBefore(Arg.getExpr()))
+    return Action::SkipChildren();
 
-ASTWalker::PostWalkResult<ArgumentList *>
-ModelASTWalker::walkToArgumentListPost(ArgumentList *ArgList) {
-  if (Expr *ParentExpr = Parent.getAsExpr()) {
-    assert(ParentExpr == ParentArgs.back().Parent &&
-           "Unmatched walkToArgumentList(Pre|Post)");
-    ParentArgs.pop_back();
+  auto *Elem = Arg.getExpr();
+  if (isa<DefaultArgumentExpr>(Elem))
+    return Action::Continue();
+
+  auto NL = Arg.getLabelLoc();
+  auto Name = Arg.getLabel();
+
+  SyntaxStructureNode SN;
+  SN.Kind = SyntaxStructureKind::Argument;
+  SN.BodyRange = charSourceRangeFromSourceRange(SM, Elem->getSourceRange());
+  if (NL.isValid() && !Name.empty()) {
+    SN.NameRange = CharSourceRange(NL, Name.getLength());
+    SN.Range = charSourceRangeFromSourceRange(
+        SM, SourceRange(NL, Elem->getEndLoc()));
+    passTokenNodesUntil(NL, ExcludeNodeAtLocation);
+  } else {
+    SN.Range = SN.BodyRange;
   }
-  return Action::Continue(ArgList);
+
+  pushStructureNode(SN, Elem);
+  return Action::Continue();
 }
 
 ASTWalker::PreWalkResult<Expr *> ModelASTWalker::walkToExprPre(Expr *E) {
   if (isVisitedBefore(E))
     return Action::SkipChildren(E);
 
-  auto addCallArgExpr = [&](const Argument &Arg) {
-    auto *Elem = Arg.getExpr();
-    if (isa<DefaultArgumentExpr>(Elem))
-      return;
-
-    auto NL = Arg.getLabelLoc();
-    auto Name = Arg.getLabel();
-
-    SyntaxStructureNode SN;
-    SN.Kind = SyntaxStructureKind::Argument;
-    SN.BodyRange = charSourceRangeFromSourceRange(SM, Elem->getSourceRange());
-    if (NL.isValid() && !Name.empty()) {
-      SN.NameRange = CharSourceRange(NL, Name.getLength());
-      SN.Range = charSourceRangeFromSourceRange(
-          SM, SourceRange(NL, Elem->getEndLoc()));
-      passTokenNodesUntil(NL, ExcludeNodeAtLocation);
-    } else {
-      SN.Range = SN.BodyRange;
-    }
-
-    pushStructureNode(SN, Elem);
-  };
-
-  if (auto *ParentExpr = Parent.getAsExpr()) {
-    if (!ParentArgs.empty() && ParentArgs.back().Parent == ParentExpr) {
-      auto &ArgumentInfo = ParentArgs.back().Args;
-      auto Arg = ArgumentInfo.find(E);
-      if (Arg != ArgumentInfo.end())
-        addCallArgExpr(Arg->second);
-    }
-  }
-
   if (E->isImplicit())
     return Action::Continue(E);