[astgen] Introduce ASTNode; update ResultType to be ASTNode.

Author: zoecarver

lib/ASTGen/Sources/ASTGen/ASTGen.swift

@@ -17,6 +17,26 @@ extension UnsafePointer {
   }
 }
 
+enum ASTNode {
+  case decl(UnsafeMutableRawPointer)
+  case stmt(UnsafeMutableRawPointer)
+  case expr(UnsafeMutableRawPointer)
+  case type(UnsafeMutableRawPointer)
+  
+  var rawValue: UnsafeMutableRawPointer {
+    switch self {
+    case .decl(let ptr):
+      return ptr
+    case .stmt(let ptr):
+      return ptr
+    case .expr(let ptr):
+      return ptr
+    case .type(let ptr):
+      return ptr
+    }
+  }
+}
+
 /// Little utility wrapper that lets us have some mutable state within
 /// immutable structs, and is therefore pretty evil.
 @propertyWrapper
@@ -29,6 +49,8 @@ class Boxed<Value> {
 }
 
 struct ASTGenVisitor: SyntaxTransformVisitor {
+  typealias ResultType = ASTNode
+  
   let ctx: UnsafeMutableRawPointer
   let base: UnsafePointer<UInt8>
 
@@ -41,11 +63,11 @@ struct ASTGenVisitor: SyntaxTransformVisitor {
 //  }
 
   @_disfavoredOverload
-  public func visit(_ node: SourceFileSyntax) -> UnsafeMutableRawPointer {
+  public func visit(_ node: SourceFileSyntax) -> ASTNode {
     fatalError("Use other overload.")
   }
 
-  public func visitAny(_ node: Syntax) -> UnsafeMutableRawPointer {
+  public func visitAny(_ node: Syntax) -> ASTNode {
     fatalError("Not implemented.")
   }
 
@@ -55,13 +77,15 @@ struct ASTGenVisitor: SyntaxTransformVisitor {
 
     for element in node.statements {
       let swiftASTNodes = visit(element)
-      if element.item.is(StmtSyntax.self) {
-        out.append(SwiftTopLevelCodeDecl_createStmt(ctx, declContext, loc, swiftASTNodes, loc))
-      } else if element.item.is(ExprSyntax.self) {
-        out.append(SwiftTopLevelCodeDecl_createExpr(ctx, declContext, loc, swiftASTNodes, loc))
-      } else {
-        assert(element.item.is(DeclSyntax.self))
-        out.append(swiftASTNodes)
+      switch swiftASTNodes {
+      case .decl(let d):
+        out.append(d)
+      case .stmt(let s):
+        out.append(SwiftTopLevelCodeDecl_createStmt(ctx, declContext, loc, s, loc))
+      case .expr(let e):
+        out.append(SwiftTopLevelCodeDecl_createExpr(ctx, declContext, loc, e, loc))
+      case .type(_):
+        fatalError("Type should not exist at top level.")
       }
     }
 

lib/ASTGen/Sources/ASTGen/Decls.swift

@@ -4,35 +4,37 @@ import SwiftSyntax
 import CASTBridging
 
 extension ASTGenVisitor {
-  public func visit(_ node: TypealiasDeclSyntax) -> UnsafeMutableRawPointer {
+  public func visit(_ node: TypealiasDeclSyntax) -> ASTNode {
     let aliasLoc = self.base.advanced(by: node.typealiasKeyword.position.utf8Offset).raw
     let equalLoc = self.base.advanced(by: node.initializer.equal.position.utf8Offset).raw
     var nameText = node.identifier.text
     let name = nameText.withUTF8 { buf in
       return SwiftASTContext_getIdentifier(ctx, buf.baseAddress, buf.count)
     }
     let nameLoc = self.base.advanced(by: node.identifier.position.utf8Offset).raw
-    let genericParams = node.genericParameterClause.map(self.visit)
+    let genericParams = node.genericParameterClause.map(self.visit).map { $0.rawValue }
     let out = TypeAliasDecl_create(self.ctx, self.declContext, aliasLoc, equalLoc, name, nameLoc, genericParams)
 
     let oldDeclContext = declContext
     declContext = out.declContext
     defer { declContext = oldDeclContext }
 
-    let underlying = self.visit(node.initializer.value)
+    let underlying = self.visit(node.initializer.value).rawValue
     TypeAliasDecl_setUnderlyingTypeRepr(out.nominalDecl, underlying)
 
-    return out.decl
+    return .decl(out.decl)
   }
 
-  public func visit(_ node: StructDeclSyntax) -> UnsafeMutableRawPointer {
+  public func visit(_ node: StructDeclSyntax) -> ASTNode {
     let loc = self.base.advanced(by: node.position.utf8Offset).raw
     var nameText = node.identifier.text
     let name = nameText.withUTF8 { buf in
       return SwiftASTContext_getIdentifier(ctx, buf.baseAddress, buf.count)
     }
 
-    let genericParams = node.genericParameterClause.map(self.visit)
+    let genericParams = node.genericParameterClause
+      .map(self.visit)
+      .map { $0.rawValue }
     let out = StructDecl_create(ctx, loc, name, loc, genericParams, declContext)
     let oldDeclContext = declContext
     declContext = out.declContext
@@ -42,10 +44,10 @@ extension ASTGenVisitor {
       NominalTypeDecl_setMembers(out.nominalDecl, ref)
     }
 
-    return out.decl
+    return .decl(out.decl)
   }
 
-  public func visit(_ node: ClassDeclSyntax) -> UnsafeMutableRawPointer {
+  public func visit(_ node: ClassDeclSyntax) -> ASTNode {
     let loc = self.base.advanced(by: node.position.utf8Offset).raw
     var nameText = node.identifier.text
     let name = nameText.withUTF8 { buf in
@@ -61,31 +63,22 @@ extension ASTGenVisitor {
       NominalTypeDecl_setMembers(out.nominalDecl, ref)
     }
 
-    return out.decl
+    return .decl(out.decl)
   }
 
-  public func visit(_ node: VariableDeclSyntax) -> UnsafeMutableRawPointer {
-    let pattern = visit(node.bindings.first!.pattern)
-    let initializer = visit(node.bindings.first!.initializer!)
+  public func visit(_ node: VariableDeclSyntax) -> ASTNode {
+    let pattern = visit(node.bindings.first!.pattern).rawValue
+    let initializer = visit(node.bindings.first!.initializer!).rawValue
 
     let loc = self.base.advanced(by: node.position.utf8Offset).raw
     let isStateic = false // TODO: compute this
     let isLet = node.letOrVarKeyword.tokenKind == .letKeyword
 
     // TODO: don't drop "initializer" on the floor.
-    return SwiftVarDecl_create(ctx, nil, loc, isStateic, isLet, declContext)
+    return .decl(SwiftVarDecl_create(ctx, nil, loc, isStateic, isLet, declContext))
   }
 
-  public func visit(_ node: CodeBlockSyntax) -> UnsafeMutableRawPointer {
-    let statements = node.statements.map(self.visit)
-    let loc = self.base.advanced(by: node.position.utf8Offset).raw
-
-    return statements.withBridgedArrayRef { ref in
-      BraceStmt_createStmt(ctx, loc, ref, loc)
-    }
-  }
-
-  public func visit(_ node: FunctionParameterSyntax) -> UnsafeMutableRawPointer {
+  public func visit(_ node: FunctionParameterSyntax) -> ASTNode {
     let loc = self.base.advanced(by: node.position.utf8Offset).raw
 
     let firstName: UnsafeMutableRawPointer?
@@ -109,34 +102,34 @@ extension ASTGenVisitor {
       secondName = nil
     }
 
-    return ParamDecl_create(ctx, loc, loc, firstName, loc, secondName, declContext)
+    return .decl(ParamDecl_create(ctx, loc, loc, firstName, loc, secondName, declContext))
   }
 
-  public func visit(_ node: FunctionDeclSyntax) -> UnsafeMutableRawPointer {
+  public func visit(_ node: FunctionDeclSyntax) -> ASTNode {
     let loc = self.base.advanced(by: node.position.utf8Offset).raw
 
     var nameText = node.identifier.text
     let name = nameText.withUTF8 { buf in
       return SwiftASTContext_getIdentifier(ctx, buf.baseAddress, buf.count)
     }
 
-    let body: UnsafeMutableRawPointer?
+    let body: ASTNode?
     if let nodeBody = node.body {
       body = visit(nodeBody)
     } else {
       body = nil
     }
 
-    let returnType: UnsafeMutableRawPointer?
+    let returnType: ASTNode?
     if let output = node.signature.output {
       returnType = visit(output.returnType)
     } else {
       returnType = nil
     }
 
     let params = node.signature.input.parameterList.map { visit($0) }
-    return params.withBridgedArrayRef { ref in
-      FuncDecl_create(ctx, loc, false, loc, name, loc, false, nil, false, nil, loc, ref, loc, body, returnType, declContext)
-    }
+    return .decl(params.withBridgedArrayRef { ref in
+      FuncDecl_create(ctx, loc, false, loc, name, loc, false, nil, false, nil, loc, ref, loc, body?.rawValue, returnType?.rawValue, declContext)
+    })
   }
 }

lib/ASTGen/Sources/ASTGen/Exprs.swift

@@ -4,74 +4,74 @@ import SwiftSyntax
 import CASTBridging
 
 extension ASTGenVisitor {
-  public func visit(_ node: ClosureExprSyntax) -> UnsafeMutableRawPointer {
+  public func visit(_ node: ClosureExprSyntax) -> ASTNode {
     let statements = node.statements.map(self.visit)
     let loc = self.base.advanced(by: node.position.utf8Offset).raw
 
     let body = statements.withBridgedArrayRef { ref in
         BraceStmt_createExpr(ctx, loc, ref, loc)
     }
 
-    return ClosureExpr_create(ctx, body, declContext)
+    return .expr(ClosureExpr_create(ctx, body, declContext))
   }
 
-  public func visit(_ node: FunctionCallExprSyntax) -> UnsafeMutableRawPointer {
+  public func visit(_ node: FunctionCallExprSyntax) -> ASTNode {
     // Transform the trailing closure into an argument.
     if let trailingClosure = node.trailingClosure {
       let tupleElement = TupleExprElementSyntax(label: nil, colon: nil, expression: ExprSyntax(trailingClosure), trailingComma: nil)
 
       return visit(node.addArgument(tupleElement).withTrailingClosure(nil))
     }
 
-    let args = visit(node.argumentList)
+    let args = visit(node.argumentList).rawValue
     // TODO: hack
-    let callee = visit(node.calledExpression)
+    let callee = visit(node.calledExpression).rawValue
 
-    return SwiftFunctionCallExpr_create(self.ctx, callee, args)
+    return .expr(SwiftFunctionCallExpr_create(self.ctx, callee, args))
   }
 
-  public func visit(_ node: IdentifierExprSyntax) -> UnsafeMutableRawPointer {
+  public func visit(_ node: IdentifierExprSyntax) -> ASTNode {
     let loc = self.base.advanced(by: node.position.utf8Offset).raw
 
     var text = node.identifier.text
     let id = text.withUTF8 { buf in
       return SwiftASTContext_getIdentifier(ctx, buf.baseAddress, buf.count)
     }
 
-    return SwiftIdentifierExpr_create(ctx, id, loc)
+    return .expr(SwiftIdentifierExpr_create(ctx, id, loc))
   }
 
-  public func visit(_ node: IdentifierPatternSyntax) -> UnsafeMutableRawPointer {
+  public func visit(_ node: IdentifierPatternSyntax) -> ASTNode {
     let loc = self.base.advanced(by: node.position.utf8Offset).raw
 
     var text = node.identifier.text
     let id = text.withUTF8 { buf in
       return SwiftASTContext_getIdentifier(ctx, buf.baseAddress, buf.count)
     }
 
-    return SwiftIdentifierExpr_create(ctx, id, loc)
+    return .expr(SwiftIdentifierExpr_create(ctx, id, loc))
   }
 
-  public func visit(_ node: MemberAccessExprSyntax) -> UnsafeMutableRawPointer {
+  public func visit(_ node: MemberAccessExprSyntax) -> ASTNode {
     let loc = self.base.advanced(by: node.position.utf8Offset).raw
-    let base = visit(node.base!)
+    let base = visit(node.base!).rawValue
     var nameText = node.name.text
     let name = nameText.withUTF8 { buf in
       return SwiftASTContext_getIdentifier(ctx, buf.baseAddress, buf.count)
     }
 
-    return UnresolvedDotExpr_create(ctx, base, loc, name, loc)
+    return .expr(UnresolvedDotExpr_create(ctx, base, loc, name, loc))
   }
 
-  public func visit(_ node: TupleExprElementListSyntax) -> UnsafeMutableRawPointer {
-    let elements = node.map(self.visit)
+  public func visit(_ node: TupleExprElementListSyntax) -> ASTNode {
+    let elements = node.map(self.visit).map { $0.rawValue }
 
     // TODO: find correct paren locs.
     let lParenLoc = self.base.advanced(by: node.position.utf8Offset).raw
     let rParenLoc = self.base.advanced(by: node.position.utf8Offset).raw
 
-    return elements.withBridgedArrayRef { elementsRef in
+    return .expr(elements.withBridgedArrayRef { elementsRef in
       SwiftTupleExpr_create(self.ctx, lParenLoc, elementsRef, rParenLoc)
-    }
+    })
   }
 }

lib/ASTGen/Sources/ASTGen/Generics.swift

@@ -4,24 +4,24 @@ import SwiftSyntax
 import CASTBridging
 
 extension ASTGenVisitor {
-  func visit(_ node: GenericParameterClauseSyntax) -> UnsafeMutableRawPointer {
+  func visit(_ node: GenericParameterClauseSyntax) -> ASTNode {
     let lAngleLoc = self.base.advanced(by: node.leftAngleBracket.position.utf8Offset).raw
     let whereLoc = node.genericWhereClause.map { self.base.advanced(by: $0.whereKeyword.position.utf8Offset).raw }
     let rAngleLoc = self.base.advanced(by: node.rightAngleBracket.position.utf8Offset).raw
-    return self.withBridgedParametersAndRequirements(node) { params, reqs in
+    return .decl(self.withBridgedParametersAndRequirements(node) { params, reqs in
       return GenericParamList_create(self.ctx, lAngleLoc, params, whereLoc, reqs, rAngleLoc)
-    }
+    })
   }
 
-  func visit(_ node: GenericParameterSyntax) -> UnsafeMutableRawPointer {
+  func visit(_ node: GenericParameterSyntax) -> ASTNode {
     var nodeName = node.name.text
     let name = nodeName.withUTF8 { buf in
       return SwiftASTContext_getIdentifier(ctx, buf.baseAddress, buf.count)
     }
     let nameLoc = self.base.advanced(by: node.name.position.utf8Offset).raw
     let ellipsisLoc = node.ellipsis.map { self.base.advanced(by: $0.position.utf8Offset).raw }
 
-    return GenericTypeParamDecl_create(self.ctx, self.declContext, name, nameLoc, ellipsisLoc, node.indexInParent / 2, ellipsisLoc != nil)
+    return .decl(GenericTypeParamDecl_create(self.ctx, self.declContext, name, nameLoc, ellipsisLoc, node.indexInParent / 2, ellipsisLoc != nil))
   }
 }
 
@@ -33,33 +33,33 @@ extension ASTGenVisitor {
     var params = [UnsafeMutableRawPointer]()
     var requirements = [BridgedRequirementRepr]()
     for param in node.genericParameterList {
-      let loweredParameter = self.visit(param)
+      let loweredParameter = self.visit(param).rawValue
       params.append(loweredParameter)
 
       guard let requirement = param.inheritedType else {
         continue
       }
 
       let loweredRequirement = self.visit(requirement)
-      GenericTypeParamDecl_setInheritedType(self.ctx, loweredParameter, loweredRequirement);
+      GenericTypeParamDecl_setInheritedType(self.ctx, loweredParameter, loweredRequirement.rawValue);
     }
 
     if let nodeRequirements = node.genericWhereClause?.requirementList {
       for requirement in nodeRequirements {
         switch requirement.body {
         case .conformanceRequirement(let conformance):
-          let firstType = self.visit(conformance.leftTypeIdentifier)
+          let firstType = self.visit(conformance.leftTypeIdentifier).rawValue
           let separatorLoc = self.base.advanced(by: conformance.colon.position.utf8Offset).raw
-          let secondType = self.visit(conformance.rightTypeIdentifier)
+          let secondType = self.visit(conformance.rightTypeIdentifier).rawValue
           requirements.append(BridgedRequirementRepr(
             SeparatorLoc: separatorLoc,
             Kind: .typeConstraint,
             FirstType: firstType,
             SecondType: secondType))
         case .sameTypeRequirement(let sameType):
-          let firstType = self.visit(sameType.leftTypeIdentifier)
+          let firstType = self.visit(sameType.leftTypeIdentifier).rawValue
           let separatorLoc = self.base.advanced(by: sameType.equalityToken.position.utf8Offset).raw
-          let secondType = self.visit(sameType.rightTypeIdentifier)
+          let secondType = self.visit(sameType.rightTypeIdentifier).rawValue
           requirements.append(BridgedRequirementRepr(
             SeparatorLoc: separatorLoc,
             Kind: .sameType,