AST: Refactor marking nodes that return (#5275)

AST: Refactor marking nodes that return

lib/coffeescript/nodes.js

@@ -378,9 +378,17 @@
       // Only override the component `ast*` methods as needed.
       ast(o, level) {
         var astNode;
+        // Merge `level` into `o` and perform other universal checks.
         o = this.astInitialize(o, level);
+        // Create serializable representation of this node.
         astNode = this.astNode(o);
-        return this.astAddReturns(astNode);
+        if ((this.astNode != null) && this.canBeReturned) {
+          // Mark AST nodes that correspond to expressions that (implicitly) return.
+          // We can’t do this as part of `astNode` because we need to assemble child
+          // nodes first before marking the parent being returned.
+          astNode.returns = true;
+        }
+        return astNode;
       }
 
       astInitialize(o, level) {
@@ -430,17 +438,6 @@
         return jisonLocationDataToAstLocationData(this.locationData);
       }
 
-      // Mark AST nodes that correspond to expressions that (implicitly) return.
-      astAddReturns(ast) {
-        if (ast == null) {
-          return ast;
-        }
-        if (this.canBeReturned) {
-          ast.returns = true;
-        }
-        return ast;
-      }
-
       // Determines whether an AST node needs an `ExpressionStatement` wrapper.
       // Typically matches our `isStatement()` logic but this allows overriding.
       isStatementAst(o) {

src/nodes.coffee

@@ -281,9 +281,16 @@ exports.Base = class Base
   # **WARNING: DO NOT OVERRIDE THIS METHOD IN CHILD CLASSES.**
   # Only override the component `ast*` methods as needed.
   ast: (o, level) ->
+    # Merge `level` into `o` and perform other universal checks.
     o = @astInitialize o, level
+    # Create serializable representation of this node.
     astNode = @astNode o
-    @astAddReturns astNode
+    # Mark AST nodes that correspond to expressions that (implicitly) return.
+    # We can’t do this as part of `astNode` because we need to assemble child
+    # nodes first before marking the parent being returned.
+    if @astNode? and @canBeReturned
+      Object.assign astNode, {returns: yes}
+    astNode
 
   astInitialize: (o, level) ->
     o = Object.assign {}, o
@@ -318,12 +325,6 @@ exports.Base = class Base
   astLocationData: ->
     jisonLocationDataToAstLocationData @locationData
 
-  # Mark AST nodes that correspond to expressions that (implicitly) return.
-  astAddReturns: (ast) ->
-    return ast unless ast?
-    ast.returns = yes if @canBeReturned
-    ast
-
   # Determines whether an AST node needs an `ExpressionStatement` wrapper.
   # Typically matches our `isStatement()` logic but this allows overriding.
   isStatementAst: (o) ->