better support for inlining new generics

Author: jdw170000

pylingual/editable_bytecode/EditableBytecode.py

@@ -152,6 +152,7 @@ def inline_annotate_functions(self):
             ("SET_FUNCTION_ATTRIBUTE", 8), # closure
             ("STORE_NAME", "__annotate_func__"),
         )
+
         # fmt: on
 
         def try_read_annotate_func(codeobj) -> EditableBytecode | None:
@@ -264,14 +265,23 @@ def is_annotate_func_and_get_inlinable_insts(codeobj) -> tuple[bool, list[Inst]]
             # handle function definition annotations
             if inst.opname == "LOAD_CONST":
                 is_annotate_func, inlinable_insts = is_annotate_func_and_get_inlinable_insts(inst.argval)
-                if not is_annotate_func:
+                if is_annotate_func:
+                    # replace
+                    # LOAD_CONST __annotate__
+                    # MAKE_FUNCTION
+                    inline_dict[(idx, tuple(self.instructions[idx : idx + 2]))] = inlinable_insts
+                    jump_target_mapping[inst] = inlinable_insts[0]
                     continue
 
-                # replace
-                # LOAD_CONST __annotate__
-                # MAKE_FUNCTION
-                inline_dict[(idx, tuple(self.instructions[idx : idx + 2]))] = inlinable_insts
-                jump_target_mapping[inst] = inlinable_insts[0]
+                if iscode(inst.argval) and inst.argval.co_name.startswith("<generic parameters of "):
+                    # fully inline generics
+                    # replace
+                    # LOAD_CONST <generic parameters of <function name>>
+                    # MAKE_FUNCTION
+                    generic_bc = EditableBytecode(inst.argval, self.opcode, self.version)
+                    inline_dict[(idx, tuple(self.instructions[idx : idx + 2]))] = generic_bc.instructions[:-1]
+                    jump_target_mapping[inst] = generic_bc.instructions[0]
+                    continue
 
             # handle inline variable annotations
             elif inst.opname in ("LOAD_NAME", "LOAD_DEREF") and inst.argval == "__conditional_annotations__":
@@ -811,7 +821,7 @@ def insert_insts(self, insert_dict: dict[int, list[Inst]]) -> int:
                 if inst.argval not in self.co_names:
                     self.co_names.append(inst.argval)
                 inst.arg = self.co_names.index(inst.argval)
-            elif inst.optype == "free":
+            elif inst.optype == "free" or inst.optype == "local":
                 if inst.argval not in self.co_varnames:
                     self.co_varnames.append(inst.argval)
                 inst.arg = self.co_varnames.index(inst.argval)

pylingual/masking/model_disasm.py

@@ -65,10 +65,17 @@ def create_global_masker(bytecode: EditableBytecode) -> Masker:
 
         # create names
         for name in bc_co.co_names:
-            if name in global_tab:
-                continue
-            global_tab.update({bc.resolve_namespace(name): f"<mask_{global_idx}>"})
-            global_idx += 1
+            if isinstance(name, tuple):
+                for n in name:
+                    if n in global_tab:
+                        continue
+                    global_tab.update({bc.resolve_namespace(n): f"<mask_{global_idx}>"})
+                    global_idx += 1
+            else:
+                if name in global_tab:
+                    continue
+                global_tab.update({bc.resolve_namespace(name): f"<mask_{global_idx}>"})
+                global_idx += 1
 
         for free in bc_co.co_freevars:
             if free in global_tab:
@@ -84,10 +91,17 @@ def create_global_masker(bytecode: EditableBytecode) -> Masker:
                 global_idx += 1
 
         for local in bc_co.co_varnames:
-            if local in global_tab:
-                continue
-            global_tab.update({bc.resolve_namespace(local): f"<mask_{global_idx}>"})
-            global_idx += 1
+            if isinstance(local, tuple):
+                for local_item in local:
+                    if local_item in global_tab:
+                        continue
+                    global_tab.update({bc.resolve_namespace(local_item): f"<mask_{global_idx}>"})
+                    global_idx += 1
+            else:
+                if local in global_tab:
+                    continue
+                global_tab.update({bc.resolve_namespace(local): f"<mask_{global_idx}>"})
+                global_idx += 1
 
         global_tab.update({bc_co.co_name: f"<mask_{global_idx}>"})
         global_idx += 1