refactor: make each dethunk optional.

DeThunk/src/main.py

@@ -1,11 +1,56 @@
-import sys
 import os
+import argparse
 
 
-def remove_thunks(path_to_file: str):
-    if not os.path.isfile(path_to_file):
-        print('invalid file.')
-        return
+class ProcessorOptions:
+    base_dir = str()
+
+    # functions
+    remove_constructor_thunk = bool()
+    remove_destructor_thunk = bool()
+    remove_virtual_function_thunk = bool()
+
+    # variables
+    remove_virtual_table_pointer_thunk = bool()
+    restore_static_variable = bool()
+
+    def __init__(self, args):
+        self.base_dir = args.path
+        self.remove_constructor_thunk = args.remove_constructor_thunk
+        self.remove_destructor_thunk = args.remove_destructor_thunk
+        self.remove_virtual_function_thunk = args.remove_virtual_function_thunk
+        self.remove_virtual_table_pointer_thunk = args.remove_virtual_table_pointer_thunk
+        self.restore_static_variable = args.restore_static_variable
+
+        if args.all:
+            self.set_all(True)
+
+    def set_function(self, opt: bool):
+        self.remove_constructor_thunk = opt
+        self.remove_destructor_thunk = opt
+        self.remove_virtual_function_thunk = opt
+
+    def set_variable(self, opt: bool):
+        self.remove_virtual_table_pointer_thunk = opt
+        self.restore_static_variable = opt
+
+    def set_all(self, opt: bool):
+        self.set_function(opt)
+        self.set_variable(opt)
+
+
+def remove_thunks(path_to_file: str, args: ProcessorOptions):
+    assert os.path.isfile(path_to_file)
+
+    RECORDED_THUNKS = []
+    if args.remove_constructor_thunk:
+        RECORDED_THUNKS.append('// constructor thunks')
+    if args.remove_destructor_thunk:
+        RECORDED_THUNKS.append('// destructor thunk')
+    if args.remove_virtual_table_pointer_thunk:
+        RECORDED_THUNKS.append('// vftables')
+    if args.remove_virtual_function_thunk:
+        RECORDED_THUNKS.append('// virtual function thunks')
 
     with open(path_to_file, 'r', encoding='utf-8') as file:
         # states
@@ -17,7 +62,7 @@ def remove_thunks(path_to_file: str):
         content = ''
         for line in file.readlines():
             # restore static member variable thunk:
-            if '// static variables' in line:
+            if args.restore_static_variable and '// static variables' in line:
                 in_static_variable = True
                 is_modified = True
             if in_static_variable:
@@ -28,11 +73,11 @@ def remove_thunks(path_to_file: str):
                         tmpline = content[begin_pos:] + tmpline
                         content = content[:begin_pos]
                         tmpline = tmpline.strip()
+
                     # remove parameter list (convert to static variable)
                     call_spec_pos = tmpline.rfind('()')
-                    if call_spec_pos == -1:
-                        print(path_to_file)
-                        assert False
+                    assert call_spec_pos != -1
+
                     tmpline = tmpline[:call_spec_pos] + tmpline[call_spec_pos + 2 :]
 
                     # remove reference
@@ -57,12 +102,7 @@ def remove_thunks(path_to_file: str):
                 in_useless_thunk = False
                 in_static_variable = False
                 continue  # don't add this line.
-            for token in [
-                '// virtual function thunks',
-                '// constructor thunks',
-                '// vftables',
-                '// destructor thunk',
-            ]:
+            for token in RECORDED_THUNKS:
                 if token in line:
                     in_useless_thunk = True
                     is_modified = True
@@ -77,26 +117,34 @@ def remove_thunks(path_to_file: str):
                 wfile.write(content)
 
 
-def iterate_headers(target_path: str):
-    def is_header(path: str):
+def iterate_headers(args: ProcessorOptions):
+    assert os.path.isdir(args.base_dir)
+
+    def is_cxx_header(path: str):
         return path.endswith('.h') or path.endswith('.hpp')
 
-    if not os.path.isdir(target_path):
-        print('invalid path.')
-        return
-
-    for root, dirs, files in os.walk(target_path):
+    for root, dirs, files in os.walk(args.base_dir):
         for file in files:
-            if is_header(file):
-                remove_thunks(os.path.join(root, file))
+            if is_cxx_header(file):
+                remove_thunks(os.path.join(root, file), args)
 
 
 def main():
-    if len(sys.argv) < 2:
-        print('usage: main.py <PATH>')
-        return
+    parser = argparse.ArgumentParser('dethunk')
 
-    iterate_headers(sys.argv[1])
+    parser.add_argument('path', help='Path to header project.')
+
+    parser.add_argument('--remove-constructor-thunk', action='store_true')
+    parser.add_argument('--remove-destructor-thunk', action='store_true')
+    parser.add_argument('--remove-virtual-table-pointer-thunk', action='store_true')
+    parser.add_argument('--remove-virtual-function-thunk', action='store_true')
+    parser.add_argument('--restore-static-variable', action='store_true')
+
+    parser.add_argument('--all', action='store_true', help='Apply all remove/restore options.')
+
+    args = parser.parse_args()
+
+    iterate_headers(ProcessorOptions(args))
 
     print('done.')