18241002908

Committed 04 Oct 2025 06:48AM UTC coverage: 91.677% (-0.002%) from 91.679%

Build # 18241002908

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #543 from OpShin/feat/opshin-optimizations

Enable UPLC-level optimizations in OpShin

Run Details

1977 of 2348 branches covered (84.2%)

Branch coverage included in aggregate %.

17 of 17 new or added lines in 2 files covered. (100.0%)

2 existing lines in 1 file now uncovered.

4852 of 5101 relevant lines covered (95.12%)

4.7 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

88.93

/opshin/optimize/optimize_const_folding.py

import typing
from collections import defaultdict
import logging

from ast import *
from ordered_set import OrderedSet

from pycardano import PlutusData

try:
    unparse
except NameError:
    from astunparse import unparse

from ..util import CompilingNodeTransformer, CompilingNodeVisitor, OPSHIN_LOGGER
from ..type_inference import INITIAL_SCOPE

"""
Pre-evaluates constant statements
"""

ACCEPTED_ATOMIC_TYPES = [
    int,
    str,
    bytes,
    type(None),
    bool,
]

SAFE_GLOBALS_LIST = [
    abs,
    all,
    any,
    ascii,
    bin,
    bool,
    bytes,
    bytearray,
    callable,
    chr,
    classmethod,
    compile,
    complex,
    delattr,
    dict,
    dir,
    divmod,
    enumerate,
    filter,
    float,
    format,
    frozenset,
    getattr,
    hasattr,
    hash,
    hex,
    id,
    input,
    int,
    isinstance,
    issubclass,
    iter,
    len,
    list,
    map,
    max,
    min,
    next,
    object,
    oct,
    open,
    ord,
    pow,
    print,
    property,
    range,
    repr,
    reversed,
    round,
    set,
    setattr,
    slice,
    sorted,
    staticmethod,
    str,
    sum,
    super,
    tuple,
    type,
    vars,
    zip,
]
SAFE_GLOBALS = {x.__name__: x for x in SAFE_GLOBALS_LIST}


class ShallowNameDefCollector(CompilingNodeVisitor):
    step = "Collecting occurring variable names"

    def __init__(self):
        self.vars = OrderedSet()

    def visit_Name(self, node: Name) -> None:
        if isinstance(node.ctx, Store):
            self.vars.add(node.id)

    def visit_ClassDef(self, node: ClassDef):
        self.vars.add(node.name)
        # ignore the content (i.e. attribute names) of class definitions

    def visit_FunctionDef(self, node: FunctionDef):
        self.vars.add(node.name)
        # ignore the recursive stuff


class DefinedTimesVisitor(CompilingNodeVisitor):
    step = "Collecting how often variables are written"

    def __init__(self):
        self.vars = defaultdict(int)

    def visit_For(self, node: For) -> None:
        # visit twice to have all names bumped to min 2 assignments
        self.generic_visit(node)
        self.generic_visit(node)
        return
        # TODO future items: use this together with guaranteed available
        # visit twice to have this name bumped to min 2 assignments
        self.visit(node.target)
        # visit the whole function
        self.generic_visit(node)

    def visit_While(self, node: While) -> None:
        # visit twice to have all names bumped to min 2 assignments
        self.generic_visit(node)
        self.generic_visit(node)
        return
        # TODO future items: use this together with guaranteed available

    def visit_If(self, node: If) -> None:
        # TODO future items: use this together with guaranteed available
        # visit twice to have all names bumped to min 2 assignments
        self.generic_visit(node)
        self.generic_visit(node)

    def visit_Name(self, node: Name) -> None:
        if isinstance(node.ctx, Store):
            self.vars[node.id] += 1

    def visit_ClassDef(self, node: ClassDef):
        self.vars[node.name] += 1
        # ignore the content (i.e. attribute names) of class definitions

    def visit_FunctionDef(self, node: FunctionDef):
        self.vars[node.name] += 1
        # visit arguments twice, they are generally assigned more than once
        for arg in node.args.args:
            self.vars[arg.arg] += 2
        self.generic_visit(node)

    def visit_Import(self, node: Import):
        for n in node.names:
            self.vars[n] += 1

    def visit_ImportFrom(self, node: ImportFrom):
        for n in node.names:
            self.vars[n] += 1


class OptimizeConstantFolding(CompilingNodeTransformer):
    step = "Constant folding"

    def __init__(self):
        self.scopes_visible = [
            OrderedSet(INITIAL_SCOPE.keys()).difference(SAFE_GLOBALS.keys())
        ]
        self.scopes_constants = [dict()]
        self.constants = OrderedSet()

    def enter_scope(self):
        self.scopes_visible.append(OrderedSet())
        self.scopes_constants.append(dict())

    def add_var_visible(self, var: str):
        self.scopes_visible[-1].add(var)

    def add_vars_visible(self, var: typing.Iterable[str]):
        self.scopes_visible[-1].update(var)

    def add_constant(self, var: str, value: typing.Any):
        self.scopes_constants[-1][var] = value

    def visible_vars(self):
        res_set = OrderedSet()
        for s in self.scopes_visible:
            res_set.update(s)
        return res_set

    def _constant_vars(self):
        res_d = {}
        for s in self.scopes_constants:
            res_d.update(s)
        return res_d

    def exit_scope(self):
        self.scopes_visible.pop(-1)
        self.scopes_constants.pop(-1)

    def _non_overwritten_globals(self):
        overwritten_vars = self.visible_vars()

        def err():
            raise ValueError("Was overwritten!")

        non_overwritten_globals = {
            k: (v if k not in overwritten_vars else err)
            for k, v in SAFE_GLOBALS.items()
        }
        return non_overwritten_globals

    def update_constants(self, node):
        a = self._non_overwritten_globals()
        a.update(self._constant_vars())
        g = a
        l = {}
        try:
            exec(unparse(node), g, l)
        except Exception as e:
            OPSHIN_LOGGER.debug(e)
        else:
            # the class is defined and added to the globals
            self.scopes_constants[-1].update(l)

    def visit_Module(self, node: Module) -> Module:
        self.enter_scope()
        def_vars_collector = ShallowNameDefCollector()
        def_vars_collector.visit(node)
        def_vars = def_vars_collector.vars
        self.add_vars_visible(def_vars)

        constant_collector = DefinedTimesVisitor()
        constant_collector.visit(node)
        constants = constant_collector.vars
        # if it is only assigned exactly once, it must be a constant (due to immutability)
        self.constants = {c for c, i in constants.items() if i == 1}

        res = self.generic_visit(node)
        self.exit_scope()
        return res

    def visit_FunctionDef(self, node: FunctionDef) -> FunctionDef:
        self.add_var_visible(node.name)
        if node.name in self.constants:
            a = self._non_overwritten_globals()
            a.update(self._constant_vars())
            g = a
            try:
                # we need to pass the global dict as local dict here to make closures possible (rec functions)
                exec(unparse(node), g, g)
            except Exception as e:
                OPSHIN_LOGGER.debug(e)
            else:
                # the class is defined and added to the globals
                self.scopes_constants[-1][node.name] = g[node.name]

        self.enter_scope()
        self.add_vars_visible(arg.arg for arg in node.args.args)
        def_vars_collector = ShallowNameDefCollector()
        for s in node.body:
            def_vars_collector.visit(s)
        def_vars = def_vars_collector.vars
        self.add_vars_visible(def_vars)

        res_node = self.generic_visit(node)
        self.exit_scope()
        return res_node

    def visit_ClassDef(self, node: ClassDef):
        if node.name in self.constants:
            self.update_constants(node)
        return node

    def visit_ImportFrom(self, node: ImportFrom):
        if all(n in self.constants for n in node.names):
            self.update_constants(node)
        return node

    def visit_Import(self, node: Import):
        if all(n in self.constants for n in node.names):
            self.update_constants(node)
        return node

    def visit_Assign(self, node: Assign):
        if len(node.targets) != 1:
            return node
        target = node.targets[0]
        if not isinstance(target, Name):
            return node

        if target.id in self.constants:
            self.update_constants(node)
        node.value = self.visit(node.value)
        return node

    def visit_AnnAssign(self, node: AnnAssign):
        target = node.target
        if not isinstance(target, Name):
            return node

        if target.id in self.constants:
            self.update_constants(node)
        node.value = self.visit(node.value)
        return node

    def generic_visit(self, node: AST):
        node = super().generic_visit(node)
        if not isinstance(node, expr):
            # only evaluate expressions, not statements
            return node
        if isinstance(node, Constant):
            # prevents unnecessary computations
            return node
        try:
            node_source = unparse(node)
        except Exception as e:
            OPSHIN_LOGGER.debug("Error when trying to unparse node: %s", e)
            return node
        if "print(" in node_source:
            # do not optimize away print statements
            return node
        try:
            # we add preceding constant plutusdata definitions here!
            g = self._non_overwritten_globals()
            l = self._constant_vars()
            node_eval = eval(node_source, g, l)
        except Exception as e:
            OPSHIN_LOGGER.debug("Error trying to evaluate node: %s", e)
            return node

        if any(
            isinstance(node_eval, t)
            for t in ACCEPTED_ATOMIC_TYPES + [list, dict, PlutusData]
        ) and not (node_eval == [] or node_eval == {}):
            new_node = Constant(node_eval, None)
            copy_location(new_node, node)
            return new_node
        return node

1	import typing	5✔
2	from collections import defaultdict	5✔
3	import logging	5✔
4
5	from ast import *	5✔
6	from ordered_set import OrderedSet	5✔
7
8	from pycardano import PlutusData	5✔
9
10	try:	5✔
11	unparse	5✔
12	except NameError:	×
13	from astunparse import unparse	×
14
15	from ..util import CompilingNodeTransformer, CompilingNodeVisitor, OPSHIN_LOGGER	5✔
16	from ..type_inference import INITIAL_SCOPE	5✔
17
18	"""	4✔
19	Pre-evaluates constant statements
20	"""
21
22	ACCEPTED_ATOMIC_TYPES = [	5✔
23	int,
24	str,
25	bytes,
26	type(None),
27	bool,
28	]
29
30	SAFE_GLOBALS_LIST = [	5✔
31	abs,
32	all,
33	any,
34	ascii,
35	bin,
36	bool,
37	bytes,
38	bytearray,
39	callable,
40	chr,
41	classmethod,
42	compile,
43	complex,
44	delattr,
45	dict,
46	dir,
47	divmod,
48	enumerate,
49	filter,
50	float,
51	format,
52	frozenset,
53	getattr,
54	hasattr,
55	hash,
56	hex,
57	id,
58	input,
59	int,
60	isinstance,
61	issubclass,
62	iter,
63	len,
64	list,
65	map,
66	max,
67	min,
68	next,
69	object,
70	oct,
71	open,
72	ord,
73	pow,
74	print,
75	property,
76	range,
77	repr,
78	reversed,
79	round,
80	set,
81	setattr,
82	slice,
83	sorted,
84	staticmethod,
85	str,
86	sum,
87	super,
88	tuple,
89	type,
90	vars,
91	zip,
92	]
93	SAFE_GLOBALS = {x.__name__: x for x in SAFE_GLOBALS_LIST}	5!
94
95
96	class ShallowNameDefCollector(CompilingNodeVisitor):	5✔
97	step = "Collecting occurring variable names"	5✔
98
99	def __init__(self):	5✔
100	self.vars = OrderedSet()	5✔
101
102	def visit_Name(self, node: Name) -> None:	5✔
103	if isinstance(node.ctx, Store):	5✔
104	self.vars.add(node.id)	5✔
105
106	def visit_ClassDef(self, node: ClassDef):	5✔
107	self.vars.add(node.name)	5✔
108	# ignore the content (i.e. attribute names) of class definitions
109
110	def visit_FunctionDef(self, node: FunctionDef):	5✔
111	self.vars.add(node.name)	5✔
112	# ignore the recursive stuff
113
114
115	class DefinedTimesVisitor(CompilingNodeVisitor):	5✔
116	step = "Collecting how often variables are written"	5✔
117
118	def __init__(self):	5✔
119	self.vars = defaultdict(int)	5✔
120
121	def visit_For(self, node: For) -> None:	5✔
122	# visit twice to have all names bumped to min 2 assignments
123	self.generic_visit(node)	5✔
124	self.generic_visit(node)	5✔
125	return	5✔
126	# TODO future items: use this together with guaranteed available
127	# visit twice to have this name bumped to min 2 assignments
128	self.visit(node.target)
129	# visit the whole function
130	self.generic_visit(node)
131
132	def visit_While(self, node: While) -> None:	5✔
133	# visit twice to have all names bumped to min 2 assignments
134	self.generic_visit(node)	5✔
135	self.generic_visit(node)	5✔
136	return	5✔
137	# TODO future items: use this together with guaranteed available
138
139	def visit_If(self, node: If) -> None:	5✔
140	# TODO future items: use this together with guaranteed available
141	# visit twice to have all names bumped to min 2 assignments
142	self.generic_visit(node)	5✔
143	self.generic_visit(node)	5✔
144
145	def visit_Name(self, node: Name) -> None:	5✔
146	if isinstance(node.ctx, Store):	5✔
147	self.vars[node.id] += 1	5✔
148
149	def visit_ClassDef(self, node: ClassDef):	5✔
150	self.vars[node.name] += 1	5✔
151	# ignore the content (i.e. attribute names) of class definitions
152
153	def visit_FunctionDef(self, node: FunctionDef):	5✔
154	self.vars[node.name] += 1	5✔
155	# visit arguments twice, they are generally assigned more than once
156	for arg in node.args.args:	5✔
157	self.vars[arg.arg] += 2	5✔
158	self.generic_visit(node)	5✔
159
160	def visit_Import(self, node: Import):	5✔
161	for n in node.names:	×
162	self.vars[n] += 1	×
163
164	def visit_ImportFrom(self, node: ImportFrom):	5✔
165	for n in node.names:	5✔
166	self.vars[n] += 1	5✔
167
168
169	class OptimizeConstantFolding(CompilingNodeTransformer):	5✔
170	step = "Constant folding"	5✔
171
172	def __init__(self):	5✔
173	self.scopes_visible = [	5✔
174	OrderedSet(INITIAL_SCOPE.keys()).difference(SAFE_GLOBALS.keys())
175	]
176	self.scopes_constants = [dict()]	5✔
177	self.constants = OrderedSet()	5✔
178
179	def enter_scope(self):	5✔
180	self.scopes_visible.append(OrderedSet())	5✔
181	self.scopes_constants.append(dict())	5✔
182
183	def add_var_visible(self, var: str):	5✔
184	self.scopes_visible[-1].add(var)	5✔
185
186	def add_vars_visible(self, var: typing.Iterable[str]):	5✔
187	self.scopes_visible[-1].update(var)	5✔
188
189	def add_constant(self, var: str, value: typing.Any):	5✔
190	self.scopes_constants[-1][var] = value	×
191
192	def visible_vars(self):	5✔
193	res_set = OrderedSet()	5✔
194	for s in self.scopes_visible:	5✔
195	res_set.update(s)	5✔
196	return res_set	5✔
197
198	def _constant_vars(self):	5✔
199	res_d = {}	5✔
200	for s in self.scopes_constants:	5✔
201	res_d.update(s)	5✔
202	return res_d	5✔
203
204	def exit_scope(self):	5✔
205	self.scopes_visible.pop(-1)	5✔
206	self.scopes_constants.pop(-1)	5✔
207
208	def _non_overwritten_globals(self):	5✔
209	overwritten_vars = self.visible_vars()	5✔
210
211	def err():	5✔
212	raise ValueError("Was overwritten!")	×
213
214	non_overwritten_globals = {	5!
215	k: (v if k not in overwritten_vars else err)
216	for k, v in SAFE_GLOBALS.items()
217	}
218	return non_overwritten_globals	5✔
219
220	def update_constants(self, node):	5✔
221	a = self._non_overwritten_globals()	5✔
222	a.update(self._constant_vars())	5✔
223	g = a	5✔
224	l = {}	5✔
225	try:	5✔
226	exec(unparse(node), g, l)	5✔
227	except Exception as e:	5✔
228	OPSHIN_LOGGER.debug(e)	5✔
229	else:
230	# the class is defined and added to the globals
231	self.scopes_constants[-1].update(l)	5✔
232
233	def visit_Module(self, node: Module) -> Module:	5✔
234	self.enter_scope()	5✔
235	def_vars_collector = ShallowNameDefCollector()	5✔
236	def_vars_collector.visit(node)	5✔
237	def_vars = def_vars_collector.vars	5✔
238	self.add_vars_visible(def_vars)	5✔
239
240	constant_collector = DefinedTimesVisitor()	5✔
241	constant_collector.visit(node)	5✔
242	constants = constant_collector.vars	5✔
243	# if it is only assigned exactly once, it must be a constant (due to immutability)
244	self.constants = {c for c, i in constants.items() if i == 1}	5!
245
246	res = self.generic_visit(node)	5✔
247	self.exit_scope()	5✔
248	return res	5✔
249
250	def visit_FunctionDef(self, node: FunctionDef) -> FunctionDef:	5✔
251	self.add_var_visible(node.name)	5✔
252	if node.name in self.constants:	5✔
253	a = self._non_overwritten_globals()	5✔
254	a.update(self._constant_vars())	5✔
255	g = a	5✔
256	try:	5✔
257	# we need to pass the global dict as local dict here to make closures possible (rec functions)
258	exec(unparse(node), g, g)	5✔
UNCOV 259	except Exception as e:	×
UNCOV 260	OPSHIN_LOGGER.debug(e)	×
261	else:
262	# the class is defined and added to the globals
263	self.scopes_constants[-1][node.name] = g[node.name]	5✔
264
265	self.enter_scope()	5✔
266	self.add_vars_visible(arg.arg for arg in node.args.args)	5✔
267	def_vars_collector = ShallowNameDefCollector()	5✔
268	for s in node.body:	5✔
269	def_vars_collector.visit(s)	5✔
270	def_vars = def_vars_collector.vars	5✔
271	self.add_vars_visible(def_vars)	5✔
272
273	res_node = self.generic_visit(node)	5✔
274	self.exit_scope()	5✔
275	return res_node	5✔
276
277	def visit_ClassDef(self, node: ClassDef):	5✔
278	if node.name in self.constants:	5✔
279	self.update_constants(node)	5✔
280	return node	5✔
281
282	def visit_ImportFrom(self, node: ImportFrom):	5✔
283	if all(n in self.constants for n in node.names):	5!
284	self.update_constants(node)	5✔
285	return node	5✔
286
287	def visit_Import(self, node: Import):	5✔
288	if all(n in self.constants for n in node.names):	×
289	self.update_constants(node)	×
290	return node	×
291
292	def visit_Assign(self, node: Assign):	5✔
293	if len(node.targets) != 1:	5!
294	return node	×
295	target = node.targets[0]	5✔
296	if not isinstance(target, Name):	5✔
297	return node	5✔
298
299	if target.id in self.constants:	5✔
300	self.update_constants(node)	5✔
301	node.value = self.visit(node.value)	5✔
302	return node	5✔
303
304	def visit_AnnAssign(self, node: AnnAssign):	5✔
305	target = node.target	5✔
306	if not isinstance(target, Name):	5!
307	return node	×
308
309	if target.id in self.constants:	5✔
310	self.update_constants(node)	5✔
311	node.value = self.visit(node.value)	5✔
312	return node	5✔
313
314	def generic_visit(self, node: AST):	5✔
315	node = super().generic_visit(node)	5✔
316	if not isinstance(node, expr):	5✔
317	# only evaluate expressions, not statements
318	return node	5✔
319	if isinstance(node, Constant):	5✔
320	# prevents unnecessary computations
321	return node	5✔
322	try:	5✔
323	node_source = unparse(node)	5✔
324	except Exception as e:	×
325	OPSHIN_LOGGER.debug("Error when trying to unparse node: %s", e)	×
326	return node	×
327	if "print(" in node_source:	5✔
328	# do not optimize away print statements
329	return node	5✔
330	try:	5✔
331	# we add preceding constant plutusdata definitions here!
332	g = self._non_overwritten_globals()	5✔
333	l = self._constant_vars()	5✔
334	node_eval = eval(node_source, g, l)	5✔
335	except Exception as e:	5✔
336	OPSHIN_LOGGER.debug("Error trying to evaluate node: %s", e)	5✔
337	return node	5✔
338
339	if any(	5✔
340	isinstance(node_eval, t)
341	for t in ACCEPTED_ATOMIC_TYPES + [list, dict, PlutusData]
342	) and not (node_eval == [] or node_eval == {}):
343	new_node = Constant(node_eval, None)	5✔
344	copy_location(new_node, node)	5✔
345	return new_node	5✔
346	return node	5✔

OpShin / opshin / 18241002908

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous