From cac47318dcc1756770010dc57041de3229d803df Mon Sep 17 00:00:00 2001
From: maximusron <aaronjomyjoseph@gmail.com>
Date: Fri, 5 Apr 2024 16:04:16 +0200
Subject: [PATCH] Update GetFunctionNumArgs

---
 lib/Interpreter/CppInterOp.cpp                | 10 ++++----
 .../CppInterOp/FunctionReflectionTest.cpp     | 23 ++++++++++++++++++-
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/lib/Interpreter/CppInterOp.cpp b/lib/Interpreter/CppInterOp.cpp
index c138781db..d1aeb7b52 100644
--- a/lib/Interpreter/CppInterOp.cpp
+++ b/lib/Interpreter/CppInterOp.cpp
@@ -828,9 +828,12 @@ namespace Cpp {
   TCppIndex_t GetFunctionNumArgs(TCppFunction_t func)
   {
     auto *D = (clang::Decl *) func;
-    if (auto *FD = llvm::dyn_cast_or_null<FunctionDecl>(D)) {
+    if (auto *FD = llvm::dyn_cast_or_null<FunctionDecl>(D))
       return FD->getNumParams();
-    }
+
+    if (auto* FD = llvm::dyn_cast_or_null<clang::FunctionTemplateDecl>(D))
+      return (FD->getTemplatedDecl())->getNumParams();
+
     return 0;
   }
 
@@ -882,9 +885,8 @@ namespace Cpp {
   // encompassed in an anonymous namespace as follows.
   namespace {
     bool IsTemplatedFunction(Decl *D) {
-      if (llvm::isa_and_nonnull<FunctionTemplateDecl>(D)) {
+      if (llvm::isa_and_nonnull<FunctionTemplateDecl>(D))
         return true;
-      }
 
       if (auto *FD = llvm::dyn_cast_or_null<FunctionDecl>(D)) {
         auto TK = FD->getTemplatedKind();
diff --git a/unittests/CppInterOp/FunctionReflectionTest.cpp b/unittests/CppInterOp/FunctionReflectionTest.cpp
index 7b19d7c32..ea73faecd 100644
--- a/unittests/CppInterOp/FunctionReflectionTest.cpp
+++ b/unittests/CppInterOp/FunctionReflectionTest.cpp
@@ -267,21 +267,42 @@ TEST(FunctionReflectionTest, GetFunctionReturnType) {
 }
 
 TEST(FunctionReflectionTest, GetFunctionNumArgs) {
-  std::vector<Decl*> Decls, SubDecls;
+  std::vector<Decl*> Decls, TemplateSubDecls;
   std::string code = R"(
     void f1() {}
     void f2(int i, double d, long l, char ch) {}
     void f3(int i, double d, long l = 0, char ch = 'a') {}
     void f4(int i = 0, double d = 0.0, long l = 0, char ch = 'a') {}
     int a;
+
+    class MyTemplatedMethodClass {
+      template<class A>
+      char get_string(A, int i) {
+          return 'A';
+      }
+              
+      template<class A>
+      void get_size() {}
+
+      template<class A, class B>
+      long add_size (A, int i, B) {
+          return sizeof(A) + i;
+      }
+    };
+
     )";
 
   GetAllTopLevelDecls(code, Decls);
+  GetAllSubDecls(Decls[5], TemplateSubDecls);
   EXPECT_EQ(Cpp::GetFunctionNumArgs(Decls[0]), (size_t) 0);
   EXPECT_EQ(Cpp::GetFunctionNumArgs(Decls[1]), (size_t) 4);
   EXPECT_EQ(Cpp::GetFunctionNumArgs(Decls[2]), (size_t) 4);
   EXPECT_EQ(Cpp::GetFunctionNumArgs(Decls[3]), (size_t) 4);
   EXPECT_EQ(Cpp::GetFunctionNumArgs(Decls[4]), 0);
+
+  EXPECT_EQ(Cpp::GetFunctionNumArgs(TemplateSubDecls[1]), 2);
+  EXPECT_EQ(Cpp::GetFunctionNumArgs(TemplateSubDecls[2]), 0);
+  EXPECT_EQ(Cpp::GetFunctionNumArgs(TemplateSubDecls[3]), 3);
 }
 
 TEST(FunctionReflectionTest, GetFunctionRequiredArgs) {