From 5a03bb4bcbacb14c8a56db92bbf5d85dae6c2772 Mon Sep 17 00:00:00 2001
From: Phodal Huang <h@phodal.com>
Date: Wed, 12 Feb 2020 13:30:12 +0800
Subject: [PATCH] feat: <go> add function multiple returnType support

---
 .../kotlin/chapi/ast/goast/GoAstListener.kt   | 25 +++++++++++++------
 .../chapi/ast/goast/GoFullIdentListener.kt    |  2 ++
 .../ast/goast/GoFullIdentListenerTest.kt      | 15 +++++++++++
 3 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/chapi-ast-go/src/main/kotlin/chapi/ast/goast/GoAstListener.kt b/chapi-ast-go/src/main/kotlin/chapi/ast/goast/GoAstListener.kt
index 90538a1a..5cc020d9 100644
--- a/chapi-ast-go/src/main/kotlin/chapi/ast/goast/GoAstListener.kt
+++ b/chapi-ast-go/src/main/kotlin/chapi/ast/goast/GoAstListener.kt
@@ -23,16 +23,27 @@ open class GoAstListener : GoParserBaseListener() {
         return ""
     }
 
-    protected fun buildReturnTypeFromSignature(codeFunction: CodeFunction, signatureContext: GoParser.SignatureContext?): Array<CodeProperty> {
+    protected fun buildReturnTypeFromSignature(
+        codeFunction: CodeFunction,
+        signatureContext: GoParser.SignatureContext?
+    ): Array<CodeProperty> {
         val result = signatureContext!!.result()
-        var returns : Array<CodeProperty> = arrayOf()
+        var returns: Array<CodeProperty> = arrayOf()
         if (result != null) {
-            val returnType = CodeProperty(
-                TypeType = result.text,
-                TypeValue = ""
-            )
+            val parameters = result.parameters()
+            if (parameters != null) {
+                for (parameterDeclContext in parameters.parameterDecl()) {
+                    val typeType = parameterDeclContext.text
+                    val returnType = CodeProperty(TypeType = typeType, TypeValue = "")
+                    returns += returnType
+                }
+            }
 
-            returns += returnType
+            if (result.type_() != null) {
+                val typeType = result.type_().text
+                val returnType = CodeProperty(TypeType = typeType, TypeValue = "")
+                returns += returnType
+            }
         }
 
         return returns;
diff --git a/chapi-ast-go/src/main/kotlin/chapi/ast/goast/GoFullIdentListener.kt b/chapi-ast-go/src/main/kotlin/chapi/ast/goast/GoFullIdentListener.kt
index 74e0f246..52a4948f 100644
--- a/chapi-ast-go/src/main/kotlin/chapi/ast/goast/GoFullIdentListener.kt
+++ b/chapi-ast-go/src/main/kotlin/chapi/ast/goast/GoFullIdentListener.kt
@@ -45,6 +45,8 @@ class GoFullIdentListener(var fileName: String) : GoAstListener() {
             Package = codeFile.PackageName
         )
 
+        codeFunction.MultipleReturns = this.buildReturnTypeFromSignature(codeFunction, ctx.signature())
+
         defaultNode.Functions += codeFunction
     }
 
diff --git a/chapi-ast-go/src/test/kotlin/chapi/ast/goast/GoFullIdentListenerTest.kt b/chapi-ast-go/src/test/kotlin/chapi/ast/goast/GoFullIdentListenerTest.kt
index 0f2e0c03..997011ee 100644
--- a/chapi-ast-go/src/test/kotlin/chapi/ast/goast/GoFullIdentListenerTest.kt
+++ b/chapi-ast-go/src/test/kotlin/chapi/ast/goast/GoFullIdentListenerTest.kt
@@ -142,4 +142,19 @@ func add(x int, y int) int {
         assertEquals(codeFile.DataStructures[0].Functions.size, 1)
         assertEquals(codeFile.DataStructures[0].Functions[0].Name, "add")
     }
+
+    @Test
+    internal fun shouldIdentifyFunctionMultipleReturnType() {
+        var code = """
+package main
+
+func get(x int, y int) (int, int) {
+	return x, y
+}
+"""
+        val codeFile = GoAnalyser().analysis(code, "")
+        assertEquals(codeFile.DataStructures.size, 1)
+        assertEquals(codeFile.DataStructures[0].Functions.size, 1)
+        assertEquals(codeFile.DataStructures[0].Functions[0].MultipleReturns.size, 2)
+    }
 }