1#!/usr/bin/env python3
2
3import argparse
4import sys
5
6def parse_args():
7	parser = argparse.ArgumentParser(description="Align structs on single lines in BEGIN...END ALIGNED SECTION blocks")
8	parser.add_argument('file', metavar="FILE", nargs='?', help="If provided, a file to process.  Otherwise stdin")
9	return parser.parse_args()
10
11def braces_match(line):
12	open_braces = 0
13	for c in line:
14		if c == '{':
15			open_braces = open_braces + 1
16		elif c == '}':
17			if open_braces == 0:
18				sys.exit("Too many close braces on line:\n{}", line)
19			open_braces = open_braces - 1
20	return (open_braces == 0)
21
22assert(braces_match("{{{}}}"))
23assert(not braces_match("{{}{}"))
24
25def is_partial_comment(line):
26	stripped = line.strip()
27	return stripped.startswith("/*") and not stripped.endswith("*/")
28
29assert(not is_partial_comment("/* */"))
30assert(not is_partial_comment("/* /* */"))
31assert(is_partial_comment("/* */ /"))
32assert(is_partial_comment("/* */ /*"))
33
34def collineate(lines):
35	ret = []
36	current_line = ""
37	for line in lines:
38		current_line = current_line + line
39		if braces_match(current_line) and not is_partial_comment(current_line):
40			ret.append(current_line)
41			current_line = ""
42	if current_line != "":
43		ret.append(current_line)
44	return ret
45
46assert(collineate(["/*", "...", "*/"]) == ["/*...*/"])
47assert(collineate(["{", "/*", "*/", "}"]) == ["{/**/}"])
48assert(collineate(["{", "{", "}", "}"]) == ["{{}}"])
49assert(collineate(["{", "}", "{", "}"]) == ["{}", "{}"])
50
51def is_line_ignored(line):
52	stripped = line.lstrip()
53	return (stripped == "") or stripped.startswith("/*") or stripped.startswith("//")
54
55def pack(line):
56	if is_line_ignored(line):
57		return line
58
59	ret = ""
60	in_leading_whitespace = True
61	for c in line:
62		if (c == " ") or (c == "\t"):
63			if in_leading_whitespace:
64				ret = ret + c
65		else:
66			in_leading_whitespace = False
67			ret = ret + c
68	return ret
69
70def prettify(line):
71	if is_line_ignored(line):
72		return line
73
74	ret = ""
75	for c in line:
76		if (c == ","):
77			ret = ret + ", "
78		elif (c == "="):
79			ret = ret + " = "
80		elif (c == "{"):
81			ret = ret + "{ "
82		elif (c == "}"):
83			ret = ret + " }"
84		else:
85			ret = ret + c
86	return ret
87
88def has_an_offset(c):
89	return (c == "{") or (c == ",")
90
91def get_offsets(line):
92	if is_line_ignored(line):
93		return None
94	ret = []
95	offset = -1
96	for c in line:
97		offset = offset + 1
98		if has_an_offset(c):
99			ret.append(offset)
100			offset = 0
101	return ret
102
103def check_offsets(lines, offsets_list):
104	reference_index = -1
105	for line in lines:
106		reference_index = reference_index + 1
107		if not is_line_ignored(line):
108			break
109	if reference_index == len(lines):
110		return # no lines to check
111	index = -1
112	for offsets in offsets_list:
113		index = index + 1
114		if offsets == None:
115			continue
116		if len(offsets) != len(offsets_list[reference_index]):
117			sys.exit("Lines have differing numbers of offsets:\n" + lines[index] + "\n" + lines[reference_index])
118
119def collect_max_offsets(offsets_list):
120	max_offsets = None
121	for offsets in offsets_list:
122		if offsets == None:
123			continue
124		if max_offsets == None:
125			max_offsets = offsets
126			continue
127		for index in range(len(max_offsets)):
128			if max_offsets[index] < offsets[index]:
129				max_offsets[index] = offsets[index]
130	return max_offsets
131
132def fix_offsets(line, target_offsets):
133	ret = ""
134	offset_index = 0
135	offset = -1
136	for c in line:
137		offset = offset + 1
138		ret = ret + c
139		if has_an_offset(c):
140			ret = ret + " " * (target_offsets[offset_index] - offset)
141			offset_index = offset_index + 1
142			offset = 0
143	return ret
144
145def align_section(lines, args):
146	lines = collineate(lines)
147	lines = [pack(l) for l in lines]
148	lines = [prettify(l) for l in lines]
149	lines = [l.rstrip() for l in lines]
150	offsets_list = [get_offsets(l) for l in lines]
151	target_offsets = collect_max_offsets(offsets_list)
152	lines = [fix_offsets(l, target_offsets) for l in lines]
153	return lines
154
155def align_file(lines, args):
156	ret = []
157	lines_to_align = None
158
159	for line in lines:
160		if lines_to_align is None:
161			if "BEGIN ALIGNED SECTION" in line:
162				lines_to_align = []
163			ret.append(line)
164		else:
165			if "END ALIGNED SECTION" in line:
166				ret = ret + align_section(lines_to_align, args)
167				lines_to_align = None
168				ret.append(line)
169			else:
170				lines_to_align.append(line)
171
172	if lines_to_align is not None:
173		sys.exit("Incomplete aligned section")
174
175	return ret
176
177if __name__ == '__main__':
178	args = parse_args()
179	lines = []
180
181	if args.file:
182		with open(args.file) as fd:
183			lines = [l.rstrip() for l in fd.readlines()]
184	else:
185		lines = [l.rstrip() for l in sys.stdin.readlines()]
186
187	for line in align_file(lines, args):
188		print(line)
189