return self from constructor

max-leuthaeuser · max-leuthaeuser · commit abbb7ca461ee · 2025-11-12T08:14:42.000+01:00
diff --git a/joern-cli/frontends/swiftsrc2cpg/src/main/scala/io/joern/swiftsrc2cpg/astcreation/AstForDeclSyntaxCreator.scala b/joern-cli/frontends/swiftsrc2cpg/src/main/scala/io/joern/swiftsrc2cpg/astcreation/AstForDeclSyntaxCreator.scala
@@ -154,21 +154,20 @@ trait AstForDeclSyntaxCreator(implicit withSchemaValidation: ValidationMode) {
 
     methodAstParentStack.pop()
 
-    val mAst = if (methodBlockContent.isEmpty) {
-      methodStubAst(methodNode_, Seq.empty, methodReturnNode_, modifiers)
-    } else {
-      val blockNode = NewBlock()
-      localAstParentStack.push(blockNode)
-      val methodBlockContentAsts = methodBlockContent.map(m => astForDeclMember(m, typeDeclNode))
-      localAstParentStack.pop()
-      methodAstWithAnnotations(
-        methodNode_,
-        Seq.empty,
-        blockAst(blockNode, methodBlockContentAsts),
-        methodReturnNode_,
-        modifiers
-      )
-    }
+    val blockNode = NewBlock()
+    localAstParentStack.push(blockNode)
+    val selfNode = identifierNode(node, "self", "self", typeDeclNode.fullName)
+    scope.addVariableReference("self", selfNode, typeDeclNode.fullName, EvaluationStrategies.BY_REFERENCE)
+    val returnAst_             = returnAst(returnNode(node, "self"), Seq(Ast(selfNode)))
+    val methodBlockContentAsts = methodBlockContent.map(m => astForDeclMember(m, typeDeclNode)) :+ returnAst_
+    localAstParentStack.pop()
+    val mAst = methodAstWithAnnotations(
+      methodNode_,
+      Seq.empty,
+      blockAst(blockNode, methodBlockContentAsts),
+      methodReturnNode_,
+      modifiers
+    )
 
     val typeDeclAst = createFunctionTypeAndTypeDecl(methodNode_)
     Ast.storeInDiffGraph(mAst.merge(typeDeclAst), diffGraph)
@@ -702,20 +701,23 @@ trait AstForDeclSyntaxCreator(implicit withSchemaValidation: ValidationMode) {
       case Some(bodyNode: CodeBlockItemListSyntax) =>
         bodyNode.children.toList match {
           case Nil => List.empty[Ast]
-          case head :: Nil if !head.item.isInstanceOf[ReturnStmtSyntax] =>
-            val retCode = code(head)
-            List(returnAst(returnNode(head, retCode), List(astForNode(head.item))))
           case children =>
             astsForBlockElements(children)
         }
       case None =>
         List.empty[Ast]
     }
 
+    val returnAst_ = if (isConstructor(node)) {
+      val selfNode = identifierNode(node, "self", "self", parentFullName)
+      scope.addVariableReference("self", selfNode, parentFullName, EvaluationStrategies.BY_REFERENCE)
+      List(returnAst(returnNode(node, "self"), Seq(Ast(selfNode))))
+    } else List.empty
+
     val methodReturnNode_ = methodReturnNode(node, returnType)
 
     val methodBlockContentAsts = methodBlockContent.map(m => astForDeclMember(m, typeDecl.get))
-    val blockAst_              = blockAst(block, methodBlockContentAsts ++ bodyStmtAsts)
+    val blockAst_              = blockAst(block, methodBlockContentAsts ++ bodyStmtAsts ++ returnAst_)
     val astForMethod =
       methodAstWithAnnotations(
         methodNode_,
diff --git a/joern-cli/frontends/swiftsrc2cpg/src/test/scala/io/joern/swiftsrc2cpg/passes/ast/ActorTests.scala b/joern-cli/frontends/swiftsrc2cpg/src/test/scala/io/joern/swiftsrc2cpg/passes/ast/ActorTests.scala
@@ -27,6 +27,13 @@ class ActorTests extends SwiftCompilerSrc2CpgSuite {
       myActorSwiftc.fullName shouldBe "SwiftTest.MyActor1"
       myActorSwiftc.member shouldBe empty
       myActorSwiftc.boundMethod.fullName.l shouldBe List("SwiftTest.MyActor1.init:()->SwiftTest.MyActor1")
+
+      val List(constructor) = myActorSwiftc.method.isConstructor.l
+      constructor.name shouldBe "init"
+      constructor.fullName shouldBe "SwiftTest.MyActor1.init:()->SwiftTest.MyActor1"
+      val List(constructorReturnSelf) = constructor.ast.isReturn.astChildren.isIdentifier.l
+      constructorReturnSelf.name shouldBe "self"
+      constructorReturnSelf.typeFullName shouldBe "SwiftTest.MyActor1"
     }
 
     "testActor2" in {
@@ -51,6 +58,10 @@ class ActorTests extends SwiftCompilerSrc2CpgSuite {
       foo.name shouldBe "foo"
       foo.fullName shouldBe "Sources/main.swift:<global>.MyActor2.foo:(x:Swift.String)->Swift.Int"
 
+      val List(constructorReturnSelf) = constructor.ast.isReturn.astChildren.isIdentifier.l
+      constructorReturnSelf.name shouldBe "self"
+      constructorReturnSelf.typeFullName shouldBe "Sources/main.swift:<global>.MyActor2"
+
       val List(myActorSwiftc) = compilerCpg.typeDecl.nameExact("MyActor2").l
       myActorSwiftc.fullName shouldBe "SwiftTest.MyActor2"
       val List(constructorSwiftc) = myActorSwiftc.method.isConstructor.l
diff --git a/joern-cli/frontends/swiftsrc2cpg/src/test/scala/io/joern/swiftsrc2cpg/passes/ast/ClosureTests.scala b/joern-cli/frontends/swiftsrc2cpg/src/test/scala/io/joern/swiftsrc2cpg/passes/ast/ClosureTests.scala
@@ -48,7 +48,7 @@ class ClosureTests extends AstSwiftSrc2CpgSuite {
       p2.code shouldBe "s2"
       p2.name shouldBe "s2"
       p2.typeFullName shouldBe "ANY"
-      closureMethod.block.astChildren.isReturn.code.l shouldBe List("s1 > s2")
+      closureMethod.block.astChildren.code.l shouldBe List("s1 > s2")
     }
 
     "testClosure3" in {
@@ -63,7 +63,7 @@ class ClosureTests extends AstSwiftSrc2CpgSuite {
       closureMethod.code shouldBe "{ foo() }"
       closureMethod.signature shouldBe "()->ANY"
       closureMethod.parameter.size shouldBe 0
-      closureMethod.block.astChildren.isReturn.code.l shouldBe List("foo()")
+      closureMethod.block.astChildren.isCall.code.l shouldBe List("foo()")
     }
 
     "testClosure4" in {
@@ -78,7 +78,7 @@ class ClosureTests extends AstSwiftSrc2CpgSuite {
       closureMethod.code shouldBe "{ foo() }"
       closureMethod.signature shouldBe "()->ANY"
       closureMethod.parameter.size shouldBe 0
-      closureMethod.block.astChildren.isReturn.code.l shouldBe List("foo()")
+      closureMethod.block.astChildren.isCall.code.l shouldBe List("foo()")
     }
 
     "testClosure5" in {
diff --git a/joern-cli/frontends/swiftsrc2cpg/src/test/scala/io/joern/swiftsrc2cpg/passes/inheritance/ClassExtensionTests.scala b/joern-cli/frontends/swiftsrc2cpg/src/test/scala/io/joern/swiftsrc2cpg/passes/inheritance/ClassExtensionTests.scala
@@ -75,7 +75,7 @@ class ClassExtensionTests extends SwiftSrc2CpgSuite {
     typeDeclB.modifier.modifierType.l shouldBe List(ModifierTypes.PRIVATE)
     val List(bConstructor) = typeDeclB.method.isConstructor.l
     bConstructor.fullName shouldBe s"Test0.swift:<global>.B.init:()->Test0.swift:<global>.B"
-    bConstructor.block.astChildren.code.l shouldBe List("var b = 0.0")
+    bConstructor.block.astChildren.code.l shouldBe List("var b = 0.0", "self")
 
     val List(typeDeclFoo) = cpg.typeDecl.nameExact("Foo").l
     typeDeclFoo.fullName shouldBe "Test0.swift:<global>.Foo"
diff --git a/joern-cli/frontends/swiftsrc2cpg/src/test/scala/io/joern/swiftsrc2cpg/passes/inheritance/EnumerationExtensionTests.scala b/joern-cli/frontends/swiftsrc2cpg/src/test/scala/io/joern/swiftsrc2cpg/passes/inheritance/EnumerationExtensionTests.scala
@@ -77,7 +77,7 @@ class EnumerationExtensionTests extends SwiftSrc2CpgSuite {
     typeDeclB.modifier.modifierType.l shouldBe List(ModifierTypes.PRIVATE)
     val List(bConstructor) = typeDeclB.method.isConstructor.l
     bConstructor.fullName shouldBe s"Test0.swift:<global>.B.init:()->Test0.swift:<global>.B"
-    bConstructor.block.astChildren.code.l shouldBe List("var b = 0.0")
+    bConstructor.block.astChildren.code.l shouldBe List("var b = 0.0", "self")
 
     val List(typeDeclFoo) = cpg.typeDecl.nameExact("Foo").l
     typeDeclFoo.fullName shouldBe "Test0.swift:<global>.Foo"
diff --git a/joern-cli/frontends/swiftsrc2cpg/src/test/scala/io/joern/swiftsrc2cpg/passes/inheritance/ProtocolExtensionTests.scala b/joern-cli/frontends/swiftsrc2cpg/src/test/scala/io/joern/swiftsrc2cpg/passes/inheritance/ProtocolExtensionTests.scala
@@ -75,7 +75,7 @@ class ProtocolExtensionTests extends SwiftSrc2CpgSuite {
     typeDeclB.modifier.modifierType.l shouldBe List(ModifierTypes.PRIVATE)
     val List(bConstructor) = typeDeclB.method.isConstructor.l
     bConstructor.fullName shouldBe s"Test0.swift:<global>.B.init:()->Test0.swift:<global>.B"
-    bConstructor.block.astChildren.code.l shouldBe List("var b = 0.0")
+    bConstructor.block.astChildren.code.l shouldBe List("var b = 0.0", "self")
 
     val List(typeDeclFoo) = cpg.typeDecl.nameExact("Foo").l
     typeDeclFoo.fullName shouldBe "Test0.swift:<global>.Foo"
diff --git a/joern-cli/frontends/swiftsrc2cpg/src/test/scala/io/joern/swiftsrc2cpg/passes/inheritance/StructureExtensionTests.scala b/joern-cli/frontends/swiftsrc2cpg/src/test/scala/io/joern/swiftsrc2cpg/passes/inheritance/StructureExtensionTests.scala
@@ -73,7 +73,7 @@ class StructureExtensionTests extends SwiftSrc2CpgSuite {
     typeDeclB.modifier.modifierType.l shouldBe List(ModifierTypes.PRIVATE)
     val List(bConstructor) = typeDeclB.method.isConstructor.l
     bConstructor.fullName shouldBe s"Test0.swift:<global>.B.init:()->Test0.swift:<global>.B"
-    bConstructor.block.astChildren.code.l shouldBe List("var b = 0.0")
+    bConstructor.block.astChildren.code.l shouldBe List("var b = 0.0", "self")
 
     val List(typeDeclFoo) = cpg.typeDecl.nameExact("Foo").l
     typeDeclFoo.fullName shouldBe "Test0.swift:<global>.Foo"
