Added support for renaming methods using assiciated types in signature

Reviewers: michal.zielinski, jerzy.kleszcz

Reviewed By: michal.zielinski, jerzy.kleszcz

Subscribers: michal.zielinski, jerzy.kleszcz

Differential Revision: https://phabricator.polidea.com/D2869

swift/lib/Obfuscation/FunctionDeclarationParser.cpp

@@ -4,6 +4,7 @@
 #include "swift/Obfuscation/DeclarationParsingUtils.h"
 
 #include <string>
+#include <sstream>
 #include <vector>
 
 namespace swift {
@@ -21,16 +22,87 @@ llvm::Error isDeclarationSupported(const FuncDecl* Declaration) {
   }
   return llvm::Error::success();
 }
-
+
+void extractSignaturePart(const Type &Result,
+                          llvm::raw_string_ostream &StringStream,
+                          std::string Fallback) {
+
+  if (auto *ResultTuple = Result->getAs<TupleType>()) {
+    auto Results = ResultTuple->getElements();
+    if (Results.empty()) {
+      StringStream << Fallback;
+    } else {
+      StringStream << "(";
+      for (auto Argument = Results.begin();
+           Argument != Results.end();
+           ++Argument) {
+        auto Name = Argument->getName().str();
+        if (!Name.empty()) {
+          StringStream << Name << ": ";
+        }
+        if (Argument->getType()->getAs<DependentMemberType>() != nullptr) {
+          // It the signature uses the associated type, we're dropping
+          // the information about it. We don't use it's name nor information
+          // where it comes from. We do it to handle the edga case of
+          // single implementation method fulfilling two functions
+          // with associated types from two different protocols.
+          StringStream << "AssociatedType";
+        } else {
+          StringStream << Argument->getType().getString();
+        }
+        if (Argument != Results.drop_back().end()) {
+          StringStream << ", ";
+        }
+      }
+      StringStream << ")";
+    }
+  } else {
+    StringStream << Result.getString();
+  }
+}
+
+std::string extractSignature(const AbstractFunctionDecl *Declaration,
+                             std::string Fallback) {
+  if (!Declaration->hasInterfaceType()) { return Fallback; }
+
+  if (auto *InstanceFunction =
+        Declaration->getInterfaceType()->getAs<AnyFunctionType>()) {
+
+    AnyFunctionType *FunctionToParse = InstanceFunction;
+
+    if (auto *Function =
+          InstanceFunction->getResult()->getAs<AnyFunctionType>()) {
+      FunctionToParse = Function;
+    }
+
+    std::string Signature;
+    llvm::raw_string_ostream StringStream(Signature);
+    extractSignaturePart(FunctionToParse->getInput(), StringStream, "()");
+    StringStream << " -> ";
+    extractSignaturePart(FunctionToParse->getResult(), StringStream, "Void");
+    return StringStream.str();
+
+  } else {
+
+    return Fallback;
+
+  }
+}
+
 std::string functionSignature(const AbstractFunctionDecl *Declaration) {
   // The signature is available via different getters depending on whether
   // it is a method or a free function
-  std::string Interface;
+  std::string Fallback;
+
+  if (!Declaration->hasInterfaceType()) { return "no_signature"; }
+
   if (Declaration->getDeclContext()->isTypeContext()) {
-    Interface = Declaration->getMethodInterfaceType().getString();
+    Fallback = Declaration->getMethodInterfaceType().getString();
   } else {
-    Interface = Declaration->getInterfaceType().getString();
+    Fallback = Declaration->getInterfaceType().getString();
   }
+
+  auto Interface = extractSignature(Declaration, Fallback);
   return "signature." + Interface;
 }
 
@@ -72,6 +144,8 @@ functionIdentifierParts(const AbstractFunctionDecl *Declaration) {
       }
       Parts.push_back("method." + SymbolName);
     }
+
+    Parts.push_back(functionSignature(Declaration));
 
   } else {
     // This logic applies to function that
@@ -86,16 +160,16 @@ functionIdentifierParts(const AbstractFunctionDecl *Declaration) {
     //       there is no override relationship between the A.a() and B.a() in
     //       protocols. it's just a name that's the same.
     //       this simplified handling should be improved in the future.
-    ModuleNameAndParts ModuleNameAndParts;
+    ValueDecl *BaseDeclaration;
     if (SatisfiesProtocol) {
       // TODO: If the function satisfies multiple protocols, we're using
       // the module name from the first of the protocols. This may lead
       // to errors and should be changed in the future.
-      ModuleNameAndParts =
-        moduleNameAndIdentifierParts(ProtocolRequirements.front());
+      BaseDeclaration = ProtocolRequirements.front();
     } else {
-      ModuleNameAndParts = moduleNameAndIdentifierParts(ProtocolDeclaration);
+      BaseDeclaration = ProtocolDeclaration;
     }
+    auto ModuleNameAndParts = moduleNameAndIdentifierParts(BaseDeclaration);
     ModuleName = ModuleNameAndParts.first;
     Parts = ModuleNameAndParts.second;
 
@@ -104,10 +178,15 @@ functionIdentifierParts(const AbstractFunctionDecl *Declaration) {
       Parts.push_back("static");
     }
     Parts.push_back("method." + SymbolName);
+
+    if (auto *ProtocolFunctionDeclaration =
+        dyn_cast<AbstractFunctionDecl>(BaseDeclaration)) {
+      Parts.push_back(functionSignature(ProtocolFunctionDeclaration));
+    } else {
+      Parts.push_back(functionSignature(Declaration));
+    }
   }
 
-  Parts.push_back(functionSignature(Declaration));
-
   return std::make_pair(ModuleName, Parts);
 }
 

swift/test/Obfuscation/FullProcess/FunctionParameterNames/Expected/FunctionParameterNames.swift

@@ -83,7 +83,7 @@ class T1_Parent{
   init(SP1_p1: String, SP1_p2: Int) {}
   init(EP3_foo IP2_foo:Int, EP4_foo IP2_bar: T1_SampleClass) { }
   init(EP1_extp1 IP1_p1: String, EP1_extp2 IP1_p2: Int) {}
-  init(_ IP1_p1: String, EP1_extp IP1_p2: Int) {}
+  init(_ IP2_p1: String, EP1_extp IP2_p2: Int) {}
 }
 class T1_Child: T1_Parent {
   override init(SP1_p1: String, SP1_p2: Int) {
@@ -92,11 +92,11 @@ class T1_Child: T1_Parent {
   override init(EP3_foo IP3_foo:Int, EP4_foo IP3_bar: T1_SampleClass) {
     super.init(EP3_foo: IP3_foo, EP4_foo: IP3_bar)
   }
-  override init(EP1_extp1 IP2_p1: String, EP1_extp2 IP2_p2: Int) {
-    super.init(EP1_extp1: IP2_p1, EP1_extp2: IP2_p2)
+  override init(EP1_extp1 IP3_p1: String, EP1_extp2 IP3_p2: Int) {
+    super.init(EP1_extp1: IP3_p1, EP1_extp2: IP3_p2)
   }
-  override init(_ IP2_p1: String, EP1_extp IP2_p2: Int) {
-    super.init(IP2_p1, EP1_extp: IP2_p2)
+  override init(_ IP4_p1: String, EP1_extp IP4_p2: Int) {
+    super.init(IP4_p1, EP1_extp: IP4_p2)
   }
 }
 let V1_c = T1_Child(SP1_p1: "p1", SP1_p2:42)

swift/test/Obfuscation/FullProcess/Functions/Expected/Functions.swift

@@ -187,3 +187,100 @@ class T1_NestedFuncs {
     return ints
   }
 }
+
+// protocols with associated types
+struct T1_TestStruct {}
+enum T1_TestEnum {}
+
+protocol T1_ParentProtocol {
+  associatedtype Fuzz
+  associatedtype Bazz
+  associatedtype Gazz
+
+  func NF2_foo(_ IP1_indexPath: Int) -> String
+  func NF1_bar(_ IP1_fuzz: Fuzz, EP1_extBazz IP1_bazz: Bazz, EP1_extGazz IP1_gazz: Gazz, EP1_atIndexPath IP2_indexPath: Int)
+}
+
+protocol T1_ChildProtocol: T1_ParentProtocol {
+  var V1_items: [[Gazz]] { get }
+}
+
+
+protocol T1_ChildProtocol2: T1_ChildProtocol { }
+
+
+final class T1_TestClass {
+
+  var V1_items: [[Gazz]] = [[
+
+    ]]
+}
+
+extension T1_TestClass: T1_ChildProtocol2 {
+
+  func NF2_foo(_ IP1_indexPath: Int) -> String { return "" }
+
+  func NF1_bar(_ IP1_fuzz: String, EP1_extBazz IP1_bazz: T1_TestStruct, EP1_extGazz IP1_gazz: T1_TestEnum, EP1_atIndexPath IP2_indexPath: Int) {}
+}
+
+protocol T1_ParentProtocol2 {
+  associatedtype Fuzz
+  associatedtype Bazz
+  associatedtype Gazz
+
+  func NF2_bar(_ IP2_fuzz: Fuzz, EP2_extBazz IP2_bazz: Bazz, SP1_gazz: Gazz) -> Void
+}
+
+protocol T1_ParentProtocol3 {
+  associatedtype Fuzz2
+  associatedtype Bazz2
+  associatedtype Gazz2
+
+  func NF2_bar(_ IP2_fuzz: Fuzz2, EP2_extBazz IP2_bazz: Bazz2, SP1_gazz: Gazz2) -> ()
+}
+
+extension T1_TestClass: T1_ParentProtocol2, T1_ParentProtocol3 {
+  typealias Fuzz = String
+  typealias Bazz = T1_TestStruct
+  typealias Gazz = T1_TestEnum
+
+  typealias Fuzz2 = T1_TestEnum
+  typealias Bazz2 = T1_TestStruct
+  typealias Gazz2 = T1_TestEnum
+
+  func NF2_bar(_ IP2_fuzz: String, EP2_extBazz IP2_bazz: T1_TestStruct, SP1_gazz: T1_TestEnum) {}
+
+  func NF2_bar(_ IP2_fuzz: T1_TestEnum, EP2_extBazz IP2_bazz: T1_TestStruct, SP1_gazz: T1_TestEnum) {}
+}
+
+// protocol functions strikes back
+
+class T1_NotWorkingParent {
+  func NF1_addSearchItem() {
+  }
+}
+
+final class T1_NextNotWorking: T1_NotWorkingParent {
+
+  override func NF1_addSearchItem() {
+    let inserter = T1_ItemInserter()
+    do {
+      let coffee = try inserter.NF1_insertEntityWithName("")
+
+    } catch {
+
+    }
+  }
+}
+
+protocol T1_ItemInserterType {
+  associatedtype Entity
+  func NF1_insertEntityWithName(_ IP1_name: String) throws -> Entity
+}
+
+struct T1_ItemInserter: T1_ItemInserterType {
+
+  func NF1_insertEntityWithName(_ IP1_name: String) throws -> String {
+    return ""
+  }
+}

swift/test/Obfuscation/FullProcess/Functions/Functions.swift

@@ -186,5 +186,103 @@ class NestedFuncs {
     ints.append(makeInt(withIdentifier: "", model: 42))
     ints.append(makeInt(withIdentifier: "", model: 42))
     return ints
+
+  }
+}
+
+// protocols with associated types
+struct TestStruct {}
+enum TestEnum {}
+
+protocol ParentProtocol {
+  associatedtype Fuzz
+  associatedtype Bazz
+  associatedtype Gazz
+
+  func foo(_ indexPath: Int) -> String
+  func bar(_ fuzz: Fuzz, extBazz bazz: Bazz, extGazz gazz: Gazz, atIndexPath indexPath: Int)
+}
+
+protocol ChildProtocol: ParentProtocol {
+  var items: [[Gazz]] { get }
+}
+
+
+protocol ChildProtocol2: ChildProtocol { }
+
+
+final class TestClass {
+
+  var items: [[Gazz]] = [[
+
+    ]]
+}
+
+extension TestClass: ChildProtocol2 {
+
+  func foo(_ indexPath: Int) -> String { return "" }
+
+  func bar(_ fuzz: String, extBazz bazz: TestStruct, extGazz gazz: TestEnum, atIndexPath indexPath: Int) {}
+}
+
+protocol ParentProtocol2 {
+  associatedtype Fuzz
+  associatedtype Bazz
+  associatedtype Gazz
+
+  func bar(_ fuzz: Fuzz, extBazz bazz: Bazz, gazz: Gazz) -> Void
+}
+
+protocol ParentProtocol3 {
+  associatedtype Fuzz2
+  associatedtype Bazz2
+  associatedtype Gazz2
+
+  func bar(_ fuzz: Fuzz2, extBazz bazz: Bazz2, gazz: Gazz2) -> ()
+}
+
+extension TestClass: ParentProtocol2, ParentProtocol3 {
+  typealias Fuzz = String
+  typealias Bazz = TestStruct
+  typealias Gazz = TestEnum
+
+  typealias Fuzz2 = TestEnum
+  typealias Bazz2 = TestStruct
+  typealias Gazz2 = TestEnum
+
+  func bar(_ fuzz: String, extBazz bazz: TestStruct, gazz: TestEnum) {}
+
+  func bar(_ fuzz: TestEnum, extBazz bazz: TestStruct, gazz: TestEnum) {}
+}
+
+// protocol functions strikes back
+
+class NotWorkingParent {
+  func addSearchItem() {
+  }
+}
+
+final class NextNotWorking: NotWorkingParent {
+
+  override func addSearchItem() {
+    let inserter = ItemInserter()
+    do {
+      let coffee = try inserter.insertEntityWithName("")
+
+    } catch {
+
+    }
+  }
+}
+
+protocol ItemInserterType {
+  associatedtype Entity
+  func insertEntityWithName(_ name: String) throws -> Entity
+}
+
+struct ItemInserter: ItemInserterType {
+
+  func insertEntityWithName(_ name: String) throws -> String {
+    return ""
   }
 }