Swift native in/out array handling

modules/objc/generator/gen_objc.py

@@ -34,6 +34,10 @@ module_imports = []
 # the list is loaded from misc/objc/gen_dict.json defined for the module and its dependencies
 class_ignore_list = []
 
+
+# list of enum names, which should be skipped by wrapper generator
+enum_ignore_list = []
+
 # list of constant names, which should be skipped by wrapper generator
 # ignored constants can be defined using regular expressions
 const_ignore_list = []
@@ -46,16 +50,16 @@ missing_consts = {}
 
 type_dict = {
     ""        : {"objc_type" : ""}, # c-tor ret_type
-    "void"    : {"objc_type" : "void", "is_primitive" : True},
-    "bool"    : {"objc_type" : "BOOL", "is_primitive" : True, "to_cpp": "(bool)%(n)s"},
-    "char"    : {"objc_type" : "char", "is_primitive" : True},
-    "int"     : {"objc_type" : "int", "is_primitive" : True, "out_type" : "int*", "out_type_ptr": "%(n)s", "out_type_ref": "*(int*)(%(n)s)"},
-    "long"    : {"objc_type" : "long", "is_primitive" : True},
-    "float"   : {"objc_type" : "float", "is_primitive" : True, "out_type" : "float*", "out_type_ptr": "%(n)s", "out_type_ref": "*(float*)(%(n)s)"},
-    "double"  : {"objc_type" : "double", "is_primitive" : True, "out_type" : "double*", "out_type_ptr": "%(n)s", "out_type_ref": "*(double*)(%(n)s)"},
+    "void"    : {"objc_type" : "void", "is_primitive" : True, "swift_type": "Void"},
+    "bool"    : {"objc_type" : "BOOL", "is_primitive" : True, "to_cpp": "(bool)%(n)s", "swift_type": "Bool"},
+    "char"    : {"objc_type" : "char", "is_primitive" : True, "swift_type": "Int8"},
+    "int"     : {"objc_type" : "int", "is_primitive" : True, "out_type" : "int*", "out_type_ptr": "%(n)s", "out_type_ref": "*(int*)(%(n)s)", "swift_type": "Int32"},
+    "long"    : {"objc_type" : "long", "is_primitive" : True, "swift_type": "Int"},
+    "float"   : {"objc_type" : "float", "is_primitive" : True, "out_type" : "float*", "out_type_ptr": "%(n)s", "out_type_ref": "*(float*)(%(n)s)", "swift_type": "Float"},
+    "double"  : {"objc_type" : "double", "is_primitive" : True, "out_type" : "double*", "out_type_ptr": "%(n)s", "out_type_ref": "*(double*)(%(n)s)", "swift_type": "Double"},
     "size_t"  : {"objc_type" : "size_t", "is_primitive" : True},
-    "int64"   : {"objc_type" : "long", "is_primitive" : True},
-    "string"  : {"objc_type" : "NSString*", "is_primitive" : True, "from_cpp": "[NSString stringWithUTF8String:%(n)s.c_str()]", "cast_to": "std::string"}
+    "int64"   : {"objc_type" : "long", "is_primitive" : True, "swift_type": "Int"},
+    "string"  : {"objc_type" : "NSString*", "is_primitive" : True, "from_cpp": "[NSString stringWithUTF8String:%(n)s.c_str()]", "cast_to": "std::string", "swift_type": "String"}
 }
 
 # Defines a rule to add extra prefixes for names from specific namespaces.
@@ -458,6 +462,7 @@ class FuncInfo(GeneralInfo):
             func_fix_map = func_arg_fix.get(self.classname or module, {}).get(self.signature(self.args), {})
             name_fix_map = func_fix_map.get(self.name, {})
             self.objc_name = name_fix_map.get('name', self.objc_name)
+            self.swift_name = name_fix_map.get('swift_name', self.swift_name)
             for arg in self.args:
                 arg_fix_map = func_fix_map.get(arg.name, {})
                 arg.ctype = arg_fix_map.get('ctype', arg.ctype) #fixing arg type
@@ -529,6 +534,62 @@ def build_objc_method_name(args):
         objc_method_name += a.name + ":"
     return objc_method_name
 
+def get_swift_type(ctype):
+    has_swift_type = "swift_type" in type_dict[ctype]
+    swift_type = type_dict[ctype]["swift_type"] if has_swift_type else type_dict[ctype]["objc_type"]
+    if swift_type[-1:] == "*":
+        swift_type = swift_type[:-1]
+    if not has_swift_type:
+        if "v_type" in type_dict[ctype]:
+            swift_type = "[" + swift_type + "]"
+        elif "v_v_type" in type_dict[ctype]:
+            swift_type = "[[" + swift_type + "]]"
+    return swift_type
+
+def build_swift_extension_decl(name, args, constructor, static, ret_type):
+    extension_decl = ("class " if static else "") + (("func " + name) if not constructor else "convenience init") + "("
+    swift_args = []
+    for a in args:
+        if a.ctype not in type_dict:
+            if not a.defval and a.ctype.endswith("*"):
+                a.defval = 0
+            if a.defval:
+                a.ctype = ''
+                continue
+        if not a.ctype:  # hidden
+            continue
+        swift_type = get_swift_type(a.ctype)
+
+        if "O" in a.out:
+            if type_dict[a.ctype].get("primitive_type", False):
+                swift_type = "UnsafeMutablePointer<" + swift_type + ">"
+            elif "v_type" in type_dict[a.ctype] or "v_v_type" in type_dict[a.ctype] or type_dict[a.ctype].get("primitive_vector", False) or type_dict[a.ctype].get("primitive_vector_vector", False):
+                swift_type = "inout " + swift_type
+
+        swift_args.append(a.name + ': ' + swift_type)
+
+    extension_decl += ", ".join(swift_args) + ")"
+    if ret_type:
+        extension_decl += " -> " + get_swift_type(ret_type)
+    return extension_decl
+
+def extension_arg(a):
+    return a.ctype in type_dict and (type_dict[a.ctype].get("primitive_vector", False) or type_dict[a.ctype].get("primitive_vector_vector", False) or (("v_type" in type_dict[a.ctype] or "v_v_type" in type_dict[a.ctype]) and "O" in a.out))
+
+def extension_tmp_arg(a):
+    if a.ctype in type_dict:
+        if type_dict[a.ctype].get("primitive_vector", False) or type_dict[a.ctype].get("primitive_vector_vector", False):
+            return a.name + "Vector"
+        elif ("v_type" in type_dict[a.ctype] or "v_v_type" in type_dict[a.ctype]) and "O" in a.out:
+            return a.name + "Array"
+    return a.name
+
+def make_swift_extension(args):
+    for a in args:
+        if extension_arg(a):
+            return True
+    return False
+
 def build_swift_signature(args):
     swift_signature = ""
     for a in args:
@@ -543,6 +604,55 @@ def build_swift_signature(args):
         swift_signature += a.name + ":"
     return swift_signature
 
+def build_unrefined_call(name, args, constructor, static, classname, has_ret):
+    swift_refine_call = ("let ret = " if has_ret and not constructor else "") + ((classname + ".") if static else "") + (name if not constructor else "self.init")
+    call_args = []
+    for a in args:
+        if a.ctype not in type_dict:
+            if not a.defval and a.ctype.endswith("*"):
+                a.defval = 0
+            if a.defval:
+                a.ctype = ''
+                continue
+        if not a.ctype:  # hidden
+            continue
+        call_args.append(a.name + ": " + extension_tmp_arg(a))
+    swift_refine_call += "(" + ", ".join(call_args) + ")"
+    return swift_refine_call
+
+def build_swift_logues(args):
+    prologue = []
+    epilogue = []
+    for a in args:
+        if a.ctype not in type_dict:
+            if not a.defval and a.ctype.endswith("*"):
+                a.defval = 0
+            if a.defval:
+                a.ctype = ''
+                continue
+        if not a.ctype:  # hidden
+            continue
+        if a.ctype in type_dict:
+            if type_dict[a.ctype].get("primitive_vector", False):
+                prologue.append("let " + extension_tmp_arg(a) + " = " + type_dict[a.ctype]["objc_type"][:-1] + "(" + a.name + ")")
+                if "O" in a.out:
+                    unsigned = type_dict[a.ctype].get("unsigned", False)
+                    array_prop = "array" if not unsigned else "unsignedArray"
+                    epilogue.append(a.name + ".removeAll()")
+                    epilogue.append(a.name + ".append(contentsOf: " +  extension_tmp_arg(a) + "." + array_prop + ")")
+            elif type_dict[a.ctype].get("primitive_vector_vector", False):
+                if not "O" in a.out:
+                    prologue.append("let " + extension_tmp_arg(a) + " = " + a.name + ".map {" + type_dict[a.ctype]["objc_type"][:-1] + "($0) }")
+                else:
+                    prologue.append("let " + extension_tmp_arg(a) + " = NSMutableArray(array: " + a.name + ".map {" + type_dict[a.ctype]["objc_type"][:-1] + "($0) })")
+                    epilogue.append(a.name + ".removeAll()")
+                    epilogue.append(a.name + ".append(contentsOf: " + extension_tmp_arg(a) + ".map { ($.0 as! " + type_dict[a.ctype]["objc_type"][:-1] + ").array  })")
+            elif ("v_type" in type_dict[a.ctype] or "v_v_type" in type_dict[a.ctype]) and "O" in a.out:
+                prologue.append("let " +  extension_tmp_arg(a) + " = NSMutableArray(array: " + a.name + ")")
+                epilogue.append(a.name + ".removeAll()")
+                epilogue.append(a.name + ".append(contentsOf: " +  extension_tmp_arg(a) + " as! " + get_swift_type(a.ctype) + ")")
+    return prologue, epilogue
+
 def add_method_to_dict(class_name, fi):
     static = fi.static if fi.classname else True
     if not method_dict.has_key((class_name, fi.objc_name)):
@@ -576,6 +686,7 @@ class ObjectiveCWrapperGenerator(object):
         self.classes["Mat"].namespace = "cv"
         self.module = ""
         self.Module = ""
+        self.extension_implementations = None # Swift extensions implementations stream
         self.ported_func_list = []
         self.skipped_func_list = []
         self.def_args_hist = {} # { def_args_cnt : funcs_cnt }
@@ -661,6 +772,8 @@ class ObjectiveCWrapperGenerator(object):
             ctype = normalize_class_name(enumType)
             constinfo = ConstInfo(decl[3][0], namespaces=self.namespaces, enumType=enumType)
             objc_type = enumType.rsplit(".", 1)[-1]
+            if objc_type in enum_ignore_list:
+                return
             if enum_fix.has_key(constinfo.classname):
                 objc_type = enum_fix[constinfo.classname].get(objc_type, objc_type)
             import_module = constinfo.classname if constinfo.classname and constinfo.classname != objc_type else self.Module
@@ -702,6 +815,8 @@ class ObjectiveCWrapperGenerator(object):
 
     def save(self, path, buf):
         global total_files, updated_files
+        if len(buf) == 0:
+            return
         total_files += 1
         if os.path.exists(path):
             with open(path, "rt") as f:
@@ -720,6 +835,9 @@ class ObjectiveCWrapperGenerator(object):
         self.clear()
         self.module = module
         self.Module = module.capitalize()
+        extension_implementations = StringIO() # Swift extensions implementations stream
+        extension_signatures = []
+
         # TODO: support UMat versions of declarations (implement UMat-wrapper for Java)
         parser = hdr_parser.CppHeaderParser(generate_umat_decls=False)
 
@@ -756,11 +874,13 @@ class ObjectiveCWrapperGenerator(object):
         logging.info("\n\n===== Generating... =====")
         package_path = os.path.join(output_objc_path, module)
         mkdir_p(package_path)
+        extension_file = "%s/%s/%sExt.swift" % (output_objc_path, module, self.Module)
+
         for ci in self.classes.values():
             if ci.name == "Mat":
                 continue
             ci.initCodeStreams(self.Module)
-            self.gen_class(ci, self.module)
+            self.gen_class(ci, self.module, extension_implementations, extension_signatures)
             classObjcHeaderCode = ci.generateObjcHeaderCode(self.module, self.Module, ci.objc_name)
             header_file = "%s/%s/%s.h" % (output_objc_path, module, ci.objc_name)
             self.save(header_file, classObjcHeaderCode)
@@ -768,6 +888,8 @@ class ObjectiveCWrapperGenerator(object):
             classObjcBodyCode = ci.generateObjcBodyCode(self.module, self.Module)
             self.save("%s/%s/%s.mm" % (output_objc_path, module, ci.objc_name), classObjcBodyCode)
             ci.cleanupCodeStreams()
+        self.save(extension_file, extension_implementations.getvalue())
+        extension_implementations.close()
         self.save(os.path.join(output_path, module+".txt"), self.makeReport())
 
     def makeReport(self):
@@ -812,7 +934,7 @@ class ObjectiveCWrapperGenerator(object):
             epilogue.append("CV2OBJC_CUSTOM(" + vector_full_type + ", " + objc_type[:-1] + ", " + vector_name + ", " + array_name + ", " + unconv_macro + ");")
             epilogue.append("#undef " + unconv_macro)
 
-    def gen_func(self, ci, fi):
+    def gen_func(self, ci, fi, extension_implementations, extension_signatures):
         logging.info("%s", fi)
         method_declarations = ci.method_declarations
         method_implementations = ci.method_implementations
@@ -904,6 +1026,7 @@ class ObjectiveCWrapperGenerator(object):
             # calculate method signature to check for uniqueness
             objc_args = build_objc_args(args)
             objc_signature = fi.signature(args)
+            swift_ext = make_swift_extension(args)
             logging.info("Objective-C: " + objc_signature)
 
             if objc_signature in objc_signatures:
@@ -1034,6 +1157,32 @@ class ObjectiveCWrapperGenerator(object):
                     tail = tail
                 )
             )
+
+            if swift_ext:
+                prototype = build_swift_extension_decl(fi.swift_name, args, constructor, static, ret_type)
+                if not (ci.name, prototype) in extension_signatures and not (ci.base, prototype) in extension_signatures:
+                    (pro, epi) = build_swift_logues(args)
+                    extension_implementations.write( Template(
+"""public extension $classname {
+    $deprecation_decl$prototype {
+$prologue
+$unrefined_call$epilogue$ret
+    }
+}
+
+"""
+                        ).substitute(
+                            classname = ci.name,
+                            deprecation_decl = "@available(*, deprecated)\n    " if fi.deprecated else "",
+                            prototype = prototype,
+                            prologue = "        " + "\n        ".join(pro),
+                            unrefined_call = "        " + build_unrefined_call(fi.swift_name, args, constructor, static, ci.name, ret_type is not None and ret_type != "void"),
+                            epilogue = "\n        " + "\n        ".join(epi) if len(epi) > 0 else "",
+                            ret = "\n        return ret" if ret_type is not None and ret_type != "void" and not constructor else ""
+                        )
+                    )
+                extension_signatures.append((ci.name, prototype))
+
             # adding method signature to dictionary
             objc_signatures.append(objc_signature)
 
@@ -1043,7 +1192,7 @@ class ObjectiveCWrapperGenerator(object):
             else:
                 break
 
-    def gen_class(self, ci, module):
+    def gen_class(self, ci, module, extension_implementations, extension_signatures):
         logging.info("%s", ci)
         if module in AdditionalImports and (ci.name in AdditionalImports[module] or "*" in AdditionalImports[module]):
             additional_imports = []
@@ -1100,7 +1249,7 @@ typedef NS_ENUM(int, {2}) {{
 
         # methods
         for fi in ci.getAllMethods():
-            self.gen_func(ci, fi)
+            self.gen_func(ci, fi, extension_implementations, extension_signatures)
         # props
         for pi in ci.props:
             ci.method_declarations.write("\n    //\n    // C++: %s %s::%s\n    //\n\n" % (pi.ctype, ci.fullName(isCPP=True), pi.name))
@@ -1412,6 +1561,7 @@ if __name__ == "__main__":
             with open(gendict_fname) as f:
                 gen_type_dict = json.load(f)
             class_ignore_list += gen_type_dict.get("class_ignore_list", [])
+            enum_ignore_list += gen_type_dict.get("enum_ignore_list", [])
             const_ignore_list += gen_type_dict.get("const_ignore_list", [])
             const_private_list += gen_type_dict.get("const_private_list", [])
             missing_consts.update(gen_type_dict.get("missing_consts", {}))