import re
import os
from pathlib import Path
def extract_functions_and_structures(content):
items = []
lines = content.split('\n')
i = 0
while i < len(lines):
line = lines[i].strip()
if not line or line.startswith('//'):
i += 1
continue
fn_match = re.match(r'^(pub\s+)?(async\s+)?fn\s+(\w+)', line)
if fn_match:
func_name = fn_match.group(3)
start_line = i
brace_count = 0
function_lines = []
started = False
while i < len(lines):
current_line = lines[i]
function_lines.append(current_line)
for char in current_line:
if char == '{':
brace_count += 1
started = True
elif char == '}':
brace_count -= 1
i += 1
if started and brace_count == 0:
break
items.append({
'type': 'function',
'name': func_name,
'content': '\n'.join(function_lines),
'start': start_line,
'end': i
})
continue
struct_match = re.match(r'^(#\[.*\]\s*)?(pub\s+)?struct\s+(\w+)', line)
if struct_match:
struct_name = struct_match.group(3)
start_line = i
struct_lines = []
j = i - 1
while j >= 0 and (lines[j].strip().startswith('#[') or lines[j].strip() == ''):
if lines[j].strip().startswith('#['):
struct_lines.insert(0, lines[j])
j -= 1
brace_count = 0
started = False
while i < len(lines):
current_line = lines[i]
struct_lines.append(current_line)
for char in current_line:
if char == '{':
brace_count += 1
started = True
elif char == '}':
brace_count -= 1
i += 1
if started and brace_count == 0:
break
elif not started and current_line.strip().endswith(';'):
break
items.append({
'type': 'struct',
'name': struct_name,
'content': '\n'.join(struct_lines),
'start': start_line,
'end': i
})
continue
enum_match = re.match(r'^(#\[.*\]\s*)?(pub\s+)?enum\s+(\w+)', line)
if enum_match:
enum_name = enum_match.group(3)
start_line = i
enum_lines = []
j = i - 1
while j >= 0 and (lines[j].strip().startswith('#[') or lines[j].strip() == ''):
if lines[j].strip().startswith('#['):
enum_lines.insert(0, lines[j])
j -= 1
brace_count = 0
started = False
while i < len(lines):
current_line = lines[i]
enum_lines.append(current_line)
for char in current_line:
if char == '{':
brace_count += 1
started = True
elif char == '}':
brace_count -= 1
i += 1
if started and brace_count == 0:
break
items.append({
'type': 'enum',
'name': enum_name,
'content': '\n'.join(enum_lines),
'start': start_line,
'end': i
})
continue
impl_match = re.match(r'^impl\s+(\w+)', line)
if impl_match:
impl_name = impl_match.group(1)
start_line = i
impl_lines = []
brace_count = 0
started = False
while i < len(lines):
current_line = lines[i]
impl_lines.append(current_line)
for char in current_line:
if char == '{':
brace_count += 1
started = True
elif char == '}':
brace_count -= 1
i += 1
if started and brace_count == 0:
break
items.append({
'type': 'impl',
'name': f'impl_{impl_name}',
'content': '\n'.join(impl_lines),
'start': start_line,
'end': i
})
continue
i += 1
return items
def extract_imports_and_constants(content):
imports = []
constants = []
lines = content.split('\n')
i = 0
while i < len(lines):
line = lines[i].strip()
if line.startswith('use '):
imports.append(lines[i])
elif line.startswith('const '):
const_lines = [lines[i]]
i += 1
while i < len(lines) and not lines[i].strip().endswith(';'):
const_lines.append(lines[i])
i += 1
if i < len(lines):
const_lines.append(lines[i])
constants.append('\n'.join(const_lines))
continue
i += 1
return imports, constants
def main():
main_rs_path = '/home/user/aicommit/src/main.rs'
decomp_dir = '/home/user/aicommit/decomposition'
with open(main_rs_path, 'r', encoding='utf-8') as f:
content = f.read()
imports, constants = extract_imports_and_constants(content)
with open(f'{decomp_dir}/00_imports.rs', 'w', encoding='utf-8') as f:
f.write('\n'.join(imports))
with open(f'{decomp_dir}/01_constants.rs', 'w', encoding='utf-8') as f:
f.write('\n\n'.join(constants))
items = extract_functions_and_structures(content)
for idx, item in enumerate(items):
filename = f"{decomp_dir}/{idx:03d}_{item['type']}_{item['name']}.rs"
with open(filename, 'w', encoding='utf-8') as f:
f.write(item['content'])
print(f"Extracted {item['type']} '{item['name']}' to {filename}")
print(f"\nTotal items extracted: {len(items)}")
print(f"Imports saved to: {decomp_dir}/00_imports.rs")
print(f"Constants saved to: {decomp_dir}/01_constants.rs")
if __name__ == '__main__':
main()