Repository: tylerhou/fiber
Branch: main
Commit: 713939487db1
Files: 19
Total size: 106.2 KB
Directory structure:
gitextract_hzh6qio3/
├── CONTRIBUTING.md
├── LICENSE
├── README.md
├── examples/
│ ├── edit_distance.py
│ ├── fib.py
│ ├── matrix_chain_mult.py
│ ├── sum.py
│ └── tree_recursion.py
└── src/
├── expressions.py
├── fiber.py
├── fiber_test.py
├── jumps.py
├── jumps_test.py
├── manual.py
├── mappers.py
├── mappers_test.py
├── trampoline.py
├── trampoline_test.py
└── utils.py
================================================
FILE CONTENTS
================================================
================================================
FILE: CONTRIBUTING.md
================================================
# How to Contribute
We'd love to accept your patches and contributions to this project. There are
just a few small guidelines you need to follow.
## Contributor License Agreement
Contributions to this project must be accompanied by a Contributor License
Agreement. You (or your employer) retain the copyright to your contribution;
this simply gives us permission to use and redistribute your contributions as
part of the project. Head over to to see
your current agreements on file or to sign a new one.
You generally only need to submit a CLA once, so if you've already submitted one
(even if it was for a different project), you probably don't need to do it
again.
## Code Reviews
All submissions, including submissions by project members, require review. We
use GitHub pull requests for this purpose. Consult
[GitHub Help](https://help.github.com/articles/about-pull-requests/) for more
information on using pull requests.
## Community Guidelines
This project follows [Google's Open Source Community
Guidelines](https://opensource.google/conduct/).
================================================
FILE: LICENSE
================================================
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
================================================
FILE: README.md
================================================
# Fiber
Fiber implements an proof-of-concept Python decorator that rewrites a function
so that it can be paused and resumed (by moving stack variables to a heap frame
and adding if statements to simulate jumps/gotos to specific lines of code).
Then, using a trampoline function that simulates the call stack on the heap, we
can call functions that recurse arbitrarily deeply without stack overflowing
(assuming we don't run out of heap memory).
```python3
cache = {}
@fiber.fiber(locals=locals())
def fib(n):
assert n >= 0
if n in cache:
return cache[n]
if n == 0:
return 0
if n == 1:
return 1
cache[n] = fib(n-1) + fib(n-2)
return cache[n]
print(sys.getrecursionlimit()) # 1000 by default
# https://www.wolframalpha.com/input/?i=fib%281010%29+mod+10**5
print(trampoline.run(fib, [1010]) % 10 ** 5) # 74305
```
Please do not use this in production.
## TOC
* [Fiber](#fiber)
* [How it works](#how-it-works)
* [Performance](#performance)
* [Limitations](#limitations)
* [Possible improvements](#possible-improvements)
* [Questions](#questions)
* [Why didn't you use Python generators?](#why-didnt-you-use-python-generators)
* [Why did you write this?](#why-did-you-write-this)
* [Contributing](#contributing)
* [License](#license)
* [Disclaimer](#disclaimer)
## How it works
A quick refresher on the call stack: normally, when some function A calls
another function B, A is "paused" while B runs to completion. Then, once B
finishes, A is resumed.
In order to move the call stack to the heap, we need to transform function A
to (1) store all variables on the heap, and (2) be able to resume execution
at specific lines of code within the function.
The first step is easy: we rewrite all local loads and stores to instead load
and store in a frame dictionary that is passed into the function. The second is
more difficult: because Python doesn't support goto statements, we have to
insert if statements to skip the code prefix that we don't want to execute.
There are a variety of ["special
forms"](https://www.gnu.org/software/emacs/manual/html_node/elisp/Special-Forms.html)
that cannot be jumped into. These we must handle by rewriting them into a form
that we do handle.
For example, if we recursively call a function inside a for loop, we would like
to be able to resume execution on the same iteration. However, when Python
executes a for loop on an non-iterator iterable it will create a new iterator
every time. To handle this case, we rewrite for loops into the equivalent while
loop. Similarly, we must rewrite boolean expressions that short circuit (`and`,
`or`) into the equivalent if statements.
Lastly, we must replace all recursive calls and normal returns by instead
returning an instruction to a trampoline to call the child function or return
the value to the parent function, respectively.
To recap, here are the AST passes we currently implement:
1. Rewrite special forms:
- `for_to_while`: Transforms for loops into the equivalent while loops.
- `promote_while_cond`: Rewrites the while conditional to use a temporary
variable that is updated every loop iteration so that we can control when
it is evaluated (e.g. if the loop condition includes a recursive call).
- `bool_exps_to_if`: Converts `and` and `or` expressions into the
equivalent if statements.
1. `promote_to_temporary`: Assigns the results of recursive calls into
temporary variables. This is necessary when we make multiple recursive calls
in the same statement (e.g. `fib(n-1) + fib(n-2)`): we need to resume
execution in the middle of the expression.
1. `remove_trivial_temporaries`: Removes temporaries that are assigned to only
once and are directly assigned to some other variable, replacing subsequent
usages with that other variable. This helps us detect tail calls.
1. `insert_jumps`: Marks the statement after yield points (currently recursive
calls and normal returns) with a `pc` index, and inserts if statements so
that re-execution of the function will resume at that program counter.
1. `lift_locals_to_frame`: Replaces loads and stores of local variables to
loads and stores in the frame object.
1. `add_trampoline_returns`: Replaces places where we must yield (recursive
calls and normal returns) with returns to the trampoline function.
1. `fix_fn_def`: Rewrites the function defintion to take a `frame` parameter.
See the [`examples`](examples) directory for functions and the results after
each AST pass. Also, see [`src/trampoline_test.py`](src/trampoline_test.py) for
some test cases.
## Performance
A simple tail-recursive function that computes the sum of an array takes about
10-11 seconds to compute with Fiber. 1000 iterations of the equivalent for loop
takes 7-8 seconds to compute. So we are slower by roughly a factor of 1000.
```python3
lst = list(range(1, 100001))
# fiber
@fiber.fiber(locals=locals())
def sum(lst, acc):
if not lst:
return acc
return sum(lst[1:], acc + lst[0])
# for loop
total = 0
for i in lst:
total += i
print(total, trampoline.run(sum, [lst, 0])) # 5000050000, 5000050000
```
We could improve the performance of the code by eliminating redundant if
checks in the generated code. Also, as we statically know the stack variables,
we can use an array for the stack frame and integer indexes (instead of a
dictionary and string hashes + lookups). This should improve the performance
significantly, but there will still probably be a large amount of overhead.
Another performance improvement is to inline the stack array: instead of
storing a list of frames in the trampoline, we could variables directly in the
stack. Again, we can compute the frame size statically. Based on some tests in
a handwritten JavaScript implementation, this has the potential to speed up the
code by roughly a factor of 2-3, at the cost of a more complex implementation.
## Limitations
- The transformation works on the AST level, so we don't support other
decorators (for example, we cannot use
[functools.cache](https://docs.python.org/3.10/library/functools.html#functools.cache)
in the above Fibonacci example).
- The function can only access variables that are passed in the `locals=`
argument. As a consequence of this, to resolve recursive function calls,
we maintain a global mapping of all fiber functions by name. This means that
fibers must have distinct names.
- We don't support some special forms (ternaries, comprehensions). These can
easily be added as a rewrite transformation.
- We don't support exceptions. This would require us to keep track of exception
handlers in the trampoline and insert returns to the trampoline to register
and deregister handlers.
- We don't support generators. To add support, we would have to modify the
trampoline to accept another operation type (yield) that sends a value to the
function that called `next()`. Also, the trampoline would have to support
multiple call stacks.
### Possible improvements
- Improve test coverage on some of the AST transformations.
- `remove_trivial_temporaries` may have a bug if the variable that it is
replaced with is reassigned to another value.
- Support more special forms (comprehensions, generators).
- Support exceptions.
- Support recursive calls that don't read the return value.
## Questions
### Why didn't you use Python generators?
It's less interesting as the transformations are easier. Here, we are
effectively implementing generators in userspace (i.e. not needing VM support);
see the answer to the next question for why this is useful.
Also, people have used generators to do this; see [one recent generator
example](https://hurryabit.github.io/blog/stack-safety-for-free/).
### Why did you write this?
- [A+ project for CS 61A at
Berkeley.](https://web.archive.org/web/20211208153249/https://cs61a.org/articles/about/#a-grades)
During the course, we created a Scheme interpreter. The extra credit
question we to replace tail calls in Python with a return to a trampoline,
with the goal that tail call optimization in Python would let us evaluate
tail calls to arbitrary depth in Scheme, in constant space.
The test cases for the question checked whether interpreting tail-call
recursive functions in Scheme caused a Python stack overflow. Using this
Fiber implementation, (1) without tail call optimization in our trampoline,
we would still be able to pass the test cases (we just wouldn't use constant
space) and (2) we can now evaluate any Scheme expression to arbitrary depth,
even if they are not in tail form.
- The React framework has an a bug open which explores a compiler transform to
rewrite JavaScript generators to a state machine so that recursive operations
(render, reconcilation) can be written more easily. This is necessary because
some JavaScript engines still don't support generators.
This project basically implements a rough version of that compiler transform
as a proof of concept, just in Python.
https://github.com/facebook/react/pull/18942
## Contributing
See [`CONTRIBUTING.md`](CONTRIBUTING.md) for more details.
## License
Apache 2.0; see [`LICENSE`](LICENSE) for more details.
## Disclaimer
This is a personal project, not an official Google project. It is not supported
by Google and Google specifically disclaims all warranties as to its quality,
merchantability, or fitness for a particular purpose.
================================================
FILE: examples/edit_distance.py
================================================
# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Pass (0): Start
def edit_distance__helper(f, s):
if f == len(first):
return len(second) - s
if s == len(second):
return len(first) - f
if first[f] == second[s]:
return edit_distance__helper(f + 1, s + 1)
del_f = edit_distance__helper(f + 1, s) + 1
replace = edit_distance__helper(f + 1, s + 1) + 1
del_s = edit_distance__helper(f, s + 1) + 1
return min(del_f, replace, del_s)
# ------------------------------------
# Pass (1): for_to_while
def edit_distance__helper(f, s):
if f == len(first):
return len(second) - s
if s == len(second):
return len(first) - f
if first[f] == second[s]:
return edit_distance__helper(f + 1, s + 1)
del_f = edit_distance__helper(f + 1, s) + 1
replace = edit_distance__helper(f + 1, s + 1) + 1
del_s = edit_distance__helper(f, s + 1) + 1
return min(del_f, replace, del_s)
# ------------------------------------
# Pass (2): promote_while_cond
def edit_distance__helper(f, s):
if f == len(first):
return len(second) - s
if s == len(second):
return len(first) - f
if first[f] == second[s]:
return edit_distance__helper(f + 1, s + 1)
del_f = edit_distance__helper(f + 1, s) + 1
replace = edit_distance__helper(f + 1, s + 1) + 1
del_s = edit_distance__helper(f, s + 1) + 1
return min(del_f, replace, del_s)
# ------------------------------------
# Pass (3): bool_exps_to_if
def edit_distance__helper(f, s):
if f == len(first):
return len(second) - s
if s == len(second):
return len(first) - f
if first[f] == second[s]:
return edit_distance__helper(f + 1, s + 1)
del_f = edit_distance__helper(f + 1, s) + 1
replace = edit_distance__helper(f + 1, s + 1) + 1
del_s = edit_distance__helper(f, s + 1) + 1
return min(del_f, replace, del_s)
# ------------------------------------
# Pass (4): promote_to_temporary_m
def edit_distance__helper(f, s):
if f == len(first):
return len(second) - s
if s == len(second):
return len(first) - f
if first[f] == second[s]:
__tmp0__ = edit_distance__helper(f + 1, s + 1)
return __tmp0__
__tmp1__ = edit_distance__helper(f + 1, s)
del_f = __tmp1__ + 1
__tmp2__ = edit_distance__helper(f + 1, s + 1)
replace = __tmp2__ + 1
__tmp3__ = edit_distance__helper(f, s + 1)
del_s = __tmp3__ + 1
return min(del_f, replace, del_s)
# ------------------------------------
# Pass (5): remove_trivial_temporaries
def edit_distance__helper(f, s):
if f == len(first):
return len(second) - s
if s == len(second):
return len(first) - f
if first[f] == second[s]:
return edit_distance__helper(f + 1, s + 1)
__tmp1__ = edit_distance__helper(f + 1, s)
del_f = __tmp1__ + 1
__tmp2__ = edit_distance__helper(f + 1, s + 1)
replace = __tmp2__ + 1
__tmp3__ = edit_distance__helper(f, s + 1)
del_s = __tmp3__ + 1
return min(del_f, replace, del_s)
# ------------------------------------
# Pass (6): insert_jumps
def edit_distance__helper(f, s):
if __pc == 0:
if f == len(first):
if __pc == 0:
return len(second) - s
__pc = 1
if s == len(second):
if __pc == 0:
return len(first) - f
__pc = 1
if first[f] == second[s]:
if __pc == 0:
return edit_distance__helper(f + 1, s + 1)
__pc = 1
__pc = 1
if __pc == 1:
__tmp1__ = edit_distance__helper(f + 1, s)
__pc = 2
if __pc == 2:
del_f = __tmp1__ + 1
__tmp2__ = edit_distance__helper(f + 1, s + 1)
__pc = 3
if __pc == 3:
replace = __tmp2__ + 1
__tmp3__ = edit_distance__helper(f, s + 1)
__pc = 4
if __pc == 4:
del_s = __tmp3__ + 1
return min(del_f, replace, del_s)
__pc = 5
# ------------------------------------
# Pass (7): lift_locals_to_frame
def edit_distance__helper(f, s):
if frame['__pc'] == 0:
if frame['f'] == len(first):
if frame['__pc'] == 0:
return len(second) - frame['s']
frame['__pc'] = 1
if frame['s'] == len(second):
if frame['__pc'] == 0:
return len(first) - frame['f']
frame['__pc'] = 1
if first[frame['f']] == second[frame['s']]:
if frame['__pc'] == 0:
return edit_distance__helper(frame['f'] + 1, frame['s'] + 1)
frame['__pc'] = 1
frame['__pc'] = 1
if frame['__pc'] == 1:
frame['__tmp1__'] = edit_distance__helper(frame['f'] + 1, frame['s'])
frame['__pc'] = 2
if frame['__pc'] == 2:
frame['del_f'] = frame['__tmp1__'] + 1
frame['__tmp2__'] = edit_distance__helper(frame['f'] + 1, frame['s'] + 1)
frame['__pc'] = 3
if frame['__pc'] == 3:
frame['replace'] = frame['__tmp2__'] + 1
frame['__tmp3__'] = edit_distance__helper(frame['f'], frame['s'] + 1)
frame['__pc'] = 4
if frame['__pc'] == 4:
frame['del_s'] = frame['__tmp3__'] + 1
return min(frame['del_f'], frame['replace'], frame['del_s'])
frame['__pc'] = 5
# ------------------------------------
# Pass (8): add_trampoline_returns
def edit_distance__helper(f, s):
if frame['__pc'] == 0:
if frame['f'] == len(first):
if frame['__pc'] == 0:
frame['__pc'] = 1
return RetOp(value=len(second) - frame['s'])
if frame['s'] == len(second):
if frame['__pc'] == 0:
frame['__pc'] = 1
return RetOp(value=len(first) - frame['f'])
if first[frame['f']] == second[frame['s']]:
if frame['__pc'] == 0:
frame['__pc'] = 1
return TailCallOp(func='edit_distance__helper', args=[frame['f'] + 1, frame['s'] + 1], kwargs={})
frame['__pc'] = 1
if frame['__pc'] == 1:
frame['__pc'] = 2
return CallOp(func='edit_distance__helper', args=[frame['f'] + 1, frame['s']], kwargs={}, ret_variable='__tmp1__')
if frame['__pc'] == 2:
frame['del_f'] = frame['__tmp1__'] + 1
frame['__pc'] = 3
return CallOp(func='edit_distance__helper', args=[frame['f'] + 1, frame['s'] + 1], kwargs={}, ret_variable='__tmp2__')
if frame['__pc'] == 3:
frame['replace'] = frame['__tmp2__'] + 1
frame['__pc'] = 4
return CallOp(func='edit_distance__helper', args=[frame['f'], frame['s'] + 1], kwargs={}, ret_variable='__tmp3__')
if frame['__pc'] == 4:
frame['del_s'] = frame['__tmp3__'] + 1
frame['__pc'] = 5
return RetOp(value=min(frame['del_f'], frame['replace'], frame['del_s']))
# ------------------------------------
# Pass (9): fix_fn_def
def __fiberfn_edit_distance__helper(frame):
if frame['__pc'] == 0:
if frame['f'] == len(first):
if frame['__pc'] == 0:
frame['__pc'] = 1
return RetOp(value=len(second) - frame['s'])
if frame['s'] == len(second):
if frame['__pc'] == 0:
frame['__pc'] = 1
return RetOp(value=len(first) - frame['f'])
if first[frame['f']] == second[frame['s']]:
if frame['__pc'] == 0:
frame['__pc'] = 1
return TailCallOp(func='edit_distance__helper', args=[frame['f'] + 1, frame['s'] + 1], kwargs={})
frame['__pc'] = 1
if frame['__pc'] == 1:
frame['__pc'] = 2
return CallOp(func='edit_distance__helper', args=[frame['f'] + 1, frame['s']], kwargs={}, ret_variable='__tmp1__')
if frame['__pc'] == 2:
frame['del_f'] = frame['__tmp1__'] + 1
frame['__pc'] = 3
return CallOp(func='edit_distance__helper', args=[frame['f'] + 1, frame['s'] + 1], kwargs={}, ret_variable='__tmp2__')
if frame['__pc'] == 3:
frame['replace'] = frame['__tmp2__'] + 1
frame['__pc'] = 4
return CallOp(func='edit_distance__helper', args=[frame['f'], frame['s'] + 1], kwargs={}, ret_variable='__tmp3__')
if frame['__pc'] == 4:
frame['del_s'] = frame['__tmp3__'] + 1
frame['__pc'] = 5
return RetOp(value=min(frame['del_f'], frame['replace'], frame['del_s']))
# ------------------------------------
================================================
FILE: examples/fib.py
================================================
# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Pass (0): Start
def fib(n):
if n in cache:
return cache[n]
if n == 0:
return 0
if n == 1:
return 1
cache[n] = fib(n - 1) + fib(n - 2)
return cache[n]
# ------------------------------------
# Pass (1): for_to_while
def fib(n):
if n in cache:
return cache[n]
if n == 0:
return 0
if n == 1:
return 1
cache[n] = fib(n - 1) + fib(n - 2)
return cache[n]
# ------------------------------------
# Pass (2): promote_while_cond
def fib(n):
if n in cache:
return cache[n]
if n == 0:
return 0
if n == 1:
return 1
cache[n] = fib(n - 1) + fib(n - 2)
return cache[n]
# ------------------------------------
# Pass (3): bool_exps_to_if
def fib(n):
if n in cache:
return cache[n]
if n == 0:
return 0
if n == 1:
return 1
cache[n] = fib(n - 1) + fib(n - 2)
return cache[n]
# ------------------------------------
# Pass (4): promote_to_temporary_m
def fib(n):
if n in cache:
return cache[n]
if n == 0:
return 0
if n == 1:
return 1
__tmp0__ = fib(n - 1)
__tmp1__ = fib(n - 2)
cache[n] = __tmp0__ + __tmp1__
return cache[n]
# ------------------------------------
# Pass (5): remove_trivial_temporaries
def fib(n):
if n in cache:
return cache[n]
if n == 0:
return 0
if n == 1:
return 1
__tmp0__ = fib(n - 1)
__tmp1__ = fib(n - 2)
cache[n] = __tmp0__ + __tmp1__
return cache[n]
# ------------------------------------
# Pass (6): insert_jumps
def fib(n):
if __pc == 0:
if n in cache:
if __pc == 0:
return cache[n]
__pc = 1
if n == 0:
if __pc == 0:
return 0
__pc = 1
if n == 1:
if __pc == 0:
return 1
__pc = 1
__tmp0__ = fib(n - 1)
__pc = 1
if __pc == 1:
__tmp1__ = fib(n - 2)
__pc = 2
if __pc == 2:
cache[n] = __tmp0__ + __tmp1__
return cache[n]
__pc = 3
# ------------------------------------
# Pass (7): lift_locals_to_frame
def fib(n):
if frame['__pc'] == 0:
if frame['n'] in cache:
if frame['__pc'] == 0:
return cache[frame['n']]
frame['__pc'] = 1
if frame['n'] == 0:
if frame['__pc'] == 0:
return 0
frame['__pc'] = 1
if frame['n'] == 1:
if frame['__pc'] == 0:
return 1
frame['__pc'] = 1
frame['__tmp0__'] = fib(frame['n'] - 1)
frame['__pc'] = 1
if frame['__pc'] == 1:
frame['__tmp1__'] = fib(frame['n'] - 2)
frame['__pc'] = 2
if frame['__pc'] == 2:
cache[frame['n']] = frame['__tmp0__'] + frame['__tmp1__']
return cache[frame['n']]
frame['__pc'] = 3
# ------------------------------------
# Pass (8): add_trampoline_returns
def fib(n):
if frame['__pc'] == 0:
if frame['n'] in cache:
if frame['__pc'] == 0:
frame['__pc'] = 1
return RetOp(value=cache[frame['n']])
if frame['n'] == 0:
if frame['__pc'] == 0:
frame['__pc'] = 1
return RetOp(value=0)
if frame['n'] == 1:
if frame['__pc'] == 0:
frame['__pc'] = 1
return RetOp(value=1)
frame['__pc'] = 1
return CallOp(func='fib', args=[frame['n'] - 1], kwargs={}, ret_variable='__tmp0__')
if frame['__pc'] == 1:
frame['__pc'] = 2
return CallOp(func='fib', args=[frame['n'] - 2], kwargs={}, ret_variable='__tmp1__')
if frame['__pc'] == 2:
cache[frame['n']] = frame['__tmp0__'] + frame['__tmp1__']
frame['__pc'] = 3
return RetOp(value=cache[frame['n']])
# ------------------------------------
# Pass (9): fix_fn_def
def __fiberfn_fib(frame):
if frame['__pc'] == 0:
if frame['n'] in cache:
if frame['__pc'] == 0:
frame['__pc'] = 1
return RetOp(value=cache[frame['n']])
if frame['n'] == 0:
if frame['__pc'] == 0:
frame['__pc'] = 1
return RetOp(value=0)
if frame['n'] == 1:
if frame['__pc'] == 0:
frame['__pc'] = 1
return RetOp(value=1)
frame['__pc'] = 1
return CallOp(func='fib', args=[frame['n'] - 1], kwargs={}, ret_variable='__tmp0__')
if frame['__pc'] == 1:
frame['__pc'] = 2
return CallOp(func='fib', args=[frame['n'] - 2], kwargs={}, ret_variable='__tmp1__')
if frame['__pc'] == 2:
cache[frame['n']] = frame['__tmp0__'] + frame['__tmp1__']
frame['__pc'] = 3
return RetOp(value=cache[frame['n']])
# ------------------------------------
================================================
FILE: examples/matrix_chain_mult.py
================================================
# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Pass (0): Start
def matrix_chain_mult__helper(i, j):
if i == j:
return 0
max_profit = float('-inf')
for k in range(i, j):
profit = matrix_chain_mult__helper(i, k)
profit += matrix_chain_mult__helper(k + 1, j)
profit += sizes[i - 1] * sizes[k] * sizes[j]
max_profit = max(max_profit, profit)
return max_profit
# ------------------------------------
# Pass (1): for_to_while
def matrix_chain_mult__helper(i, j):
if i == j:
return 0
max_profit = float('-inf')
__tmp0__ = iter(range(i, j))
__tmp1__ = True
while __tmp1__:
try:
k = next(__tmp0__)
except StopIteration:
__tmp1__ = False
continue
profit = matrix_chain_mult__helper(i, k)
profit += matrix_chain_mult__helper(k + 1, j)
profit += sizes[i - 1] * sizes[k] * sizes[j]
max_profit = max(max_profit, profit)
return max_profit
# ------------------------------------
# Pass (2): promote_while_cond
def matrix_chain_mult__helper(i, j):
if i == j:
return 0
max_profit = float('-inf')
__tmp0__ = iter(range(i, j))
__tmp1__ = True
while __tmp1__:
try:
k = next(__tmp0__)
except StopIteration:
__tmp1__ = False
continue
profit = matrix_chain_mult__helper(i, k)
profit += matrix_chain_mult__helper(k + 1, j)
profit += sizes[i - 1] * sizes[k] * sizes[j]
max_profit = max(max_profit, profit)
return max_profit
# ------------------------------------
# Pass (3): bool_exps_to_if
def matrix_chain_mult__helper(i, j):
if i == j:
return 0
max_profit = float('-inf')
__tmp0__ = iter(range(i, j))
__tmp1__ = True
while __tmp1__:
try:
k = next(__tmp0__)
except StopIteration:
__tmp1__ = False
continue
profit = matrix_chain_mult__helper(i, k)
profit += matrix_chain_mult__helper(k + 1, j)
profit += sizes[i - 1] * sizes[k] * sizes[j]
max_profit = max(max_profit, profit)
return max_profit
# ------------------------------------
# Pass (4): promote_to_temporary_m
def matrix_chain_mult__helper(i, j):
if i == j:
return 0
max_profit = float('-inf')
__tmp0__ = iter(range(i, j))
__tmp1__ = True
while __tmp1__:
try:
k = next(__tmp0__)
except StopIteration:
__tmp1__ = False
continue
__tmp2__ = matrix_chain_mult__helper(i, k)
profit = __tmp2__
__tmp3__ = matrix_chain_mult__helper(k + 1, j)
profit += __tmp3__
profit += sizes[i - 1] * sizes[k] * sizes[j]
max_profit = max(max_profit, profit)
return max_profit
# ------------------------------------
# Pass (5): remove_trivial_temporaries
def matrix_chain_mult__helper(i, j):
if i == j:
return 0
max_profit = float('-inf')
__tmp0__ = iter(range(i, j))
__tmp1__ = True
while __tmp1__:
try:
k = next(__tmp0__)
except StopIteration:
__tmp1__ = False
continue
profit = matrix_chain_mult__helper(i, k)
__tmp3__ = matrix_chain_mult__helper(k + 1, j)
profit += __tmp3__
profit += sizes[i - 1] * sizes[k] * sizes[j]
max_profit = max(max_profit, profit)
return max_profit
# ------------------------------------
# Pass (6): insert_jumps
def matrix_chain_mult__helper(i, j):
if __pc == 0:
if i == j:
if __pc == 0:
return 0
__pc = 1
max_profit = float('-inf')
__tmp0__ = iter(range(i, j))
__tmp1__ = True
__pc = 1
if 1 <= __pc < 4:
while __tmp1__:
if __pc == 1:
try:
k = next(__tmp0__)
except StopIteration:
__tmp1__ = False
continue
profit = matrix_chain_mult__helper(i, k)
__pc = 2
if __pc == 2:
__tmp3__ = matrix_chain_mult__helper(k + 1, j)
__pc = 3
if __pc == 3:
profit += __tmp3__
profit += sizes[i - 1] * sizes[k] * sizes[j]
max_profit = max(max_profit, profit)
__pc = 4
__pc = 1
__pc = 4
if __pc == 4:
return max_profit
__pc = 5
# ------------------------------------
# Pass (7): lift_locals_to_frame
def matrix_chain_mult__helper(i, j):
if frame['__pc'] == 0:
if frame['i'] == frame['j']:
if frame['__pc'] == 0:
return 0
frame['__pc'] = 1
frame['max_profit'] = float('-inf')
frame['__tmp0__'] = iter(range(frame['i'], frame['j']))
frame['__tmp1__'] = True
frame['__pc'] = 1
if 1 <= frame['__pc'] < 4:
while frame['__tmp1__']:
if frame['__pc'] == 1:
try:
frame['k'] = next(frame['__tmp0__'])
except StopIteration:
frame['__tmp1__'] = False
continue
frame['profit'] = matrix_chain_mult__helper(frame['i'], frame['k'])
frame['__pc'] = 2
if frame['__pc'] == 2:
frame['__tmp3__'] = matrix_chain_mult__helper(frame['k'] + 1, frame['j'])
frame['__pc'] = 3
if frame['__pc'] == 3:
frame['profit'] += frame['__tmp3__']
frame['profit'] += sizes[frame['i'] - 1] * sizes[frame['k']] * sizes[frame['j']]
frame['max_profit'] = max(frame['max_profit'], frame['profit'])
frame['__pc'] = 4
frame['__pc'] = 1
frame['__pc'] = 4
if frame['__pc'] == 4:
return frame['max_profit']
frame['__pc'] = 5
# ------------------------------------
# Pass (8): add_trampoline_returns
def matrix_chain_mult__helper(i, j):
if frame['__pc'] == 0:
if frame['i'] == frame['j']:
if frame['__pc'] == 0:
frame['__pc'] = 1
return RetOp(value=0)
frame['max_profit'] = float('-inf')
frame['__tmp0__'] = iter(range(frame['i'], frame['j']))
frame['__tmp1__'] = True
frame['__pc'] = 1
if 1 <= frame['__pc'] < 4:
while frame['__tmp1__']:
if frame['__pc'] == 1:
try:
frame['k'] = next(frame['__tmp0__'])
except StopIteration:
frame['__tmp1__'] = False
continue
frame['__pc'] = 2
return CallOp(func='matrix_chain_mult__helper', args=[frame['i'], frame['k']], kwargs={}, ret_variable='profit')
if frame['__pc'] == 2:
frame['__pc'] = 3
return CallOp(func='matrix_chain_mult__helper', args=[frame['k'] + 1, frame['j']], kwargs={}, ret_variable='__tmp3__')
if frame['__pc'] == 3:
frame['profit'] += frame['__tmp3__']
frame['profit'] += sizes[frame['i'] - 1] * sizes[frame['k']] * sizes[frame['j']]
frame['max_profit'] = max(frame['max_profit'], frame['profit'])
frame['__pc'] = 4
frame['__pc'] = 1
frame['__pc'] = 4
if frame['__pc'] == 4:
frame['__pc'] = 5
return RetOp(value=frame['max_profit'])
# ------------------------------------
# Pass (9): fix_fn_def
def __fiberfn_matrix_chain_mult__helper(frame):
if frame['__pc'] == 0:
if frame['i'] == frame['j']:
if frame['__pc'] == 0:
frame['__pc'] = 1
return RetOp(value=0)
frame['max_profit'] = float('-inf')
frame['__tmp0__'] = iter(range(frame['i'], frame['j']))
frame['__tmp1__'] = True
frame['__pc'] = 1
if 1 <= frame['__pc'] < 4:
while frame['__tmp1__']:
if frame['__pc'] == 1:
try:
frame['k'] = next(frame['__tmp0__'])
except StopIteration:
frame['__tmp1__'] = False
continue
frame['__pc'] = 2
return CallOp(func='matrix_chain_mult__helper', args=[frame['i'], frame['k']], kwargs={}, ret_variable='profit')
if frame['__pc'] == 2:
frame['__pc'] = 3
return CallOp(func='matrix_chain_mult__helper', args=[frame['k'] + 1, frame['j']], kwargs={}, ret_variable='__tmp3__')
if frame['__pc'] == 3:
frame['profit'] += frame['__tmp3__']
frame['profit'] += sizes[frame['i'] - 1] * sizes[frame['k']] * sizes[frame['j']]
frame['max_profit'] = max(frame['max_profit'], frame['profit'])
frame['__pc'] = 4
frame['__pc'] = 1
frame['__pc'] = 4
if frame['__pc'] == 4:
frame['__pc'] = 5
return RetOp(value=frame['max_profit'])
# ------------------------------------
================================================
FILE: examples/sum.py
================================================
# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Pass (0): Start
def sum(lst, acc):
if not lst:
return acc
return sum(lst[1:], acc + lst[0])
# ------------------------------------
# Pass (1): for_to_while
def sum(lst, acc):
if not lst:
return acc
return sum(lst[1:], acc + lst[0])
# ------------------------------------
# Pass (2): promote_while_cond
def sum(lst, acc):
if not lst:
return acc
return sum(lst[1:], acc + lst[0])
# ------------------------------------
# Pass (3): bool_exps_to_if
def sum(lst, acc):
if not lst:
return acc
return sum(lst[1:], acc + lst[0])
# ------------------------------------
# Pass (4): promote_to_temporary_m
def sum(lst, acc):
if not lst:
return acc
__tmp0__ = sum(lst[1:], acc + lst[0])
return __tmp0__
# ------------------------------------
# Pass (5): remove_trivial_temporaries
def sum(lst, acc):
if not lst:
return acc
return sum(lst[1:], acc + lst[0])
# ------------------------------------
# Pass (6): insert_jumps
def sum(lst, acc):
if __pc == 0:
if not lst:
if __pc == 0:
return acc
__pc = 1
return sum(lst[1:], acc + lst[0])
__pc = 1
# ------------------------------------
# Pass (7): lift_locals_to_frame
def sum(lst, acc):
if frame['__pc'] == 0:
if not frame['lst']:
if frame['__pc'] == 0:
return frame['acc']
frame['__pc'] = 1
return sum(frame['lst'][1:], frame['acc'] + frame['lst'][0])
frame['__pc'] = 1
# ------------------------------------
# Pass (8): add_trampoline_returns
def sum(lst, acc):
if frame['__pc'] == 0:
if not frame['lst']:
if frame['__pc'] == 0:
frame['__pc'] = 1
return RetOp(value=frame['acc'])
frame['__pc'] = 1
return TailCallOp(func='sum', args=[frame['lst'][1:], frame['acc'] + frame['lst'][0]], kwargs={})
# ------------------------------------
# Pass (9): fix_fn_def
def __fiberfn_sum(frame):
if frame['__pc'] == 0:
if not frame['lst']:
if frame['__pc'] == 0:
frame['__pc'] = 1
return RetOp(value=frame['acc'])
frame['__pc'] = 1
return TailCallOp(func='sum', args=[frame['lst'][1:], frame['acc'] + frame['lst'][0]], kwargs={})
# ------------------------------------
================================================
FILE: examples/tree_recursion.py
================================================
# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Pass (0): Start
def all_zeroes(tree):
if tree == 0:
return True
if not isinstance(tree, Tree):
return False
return all_zeroes(tree.left) and all_zeroes(tree.right)
# ------------------------------------
# Pass (1): for_to_while
def all_zeroes(tree):
if tree == 0:
return True
if not isinstance(tree, Tree):
return False
return all_zeroes(tree.left) and all_zeroes(tree.right)
# ------------------------------------
# Pass (2): promote_while_cond
def all_zeroes(tree):
if tree == 0:
return True
if not isinstance(tree, Tree):
return False
return all_zeroes(tree.left) and all_zeroes(tree.right)
# ------------------------------------
# Pass (3): bool_exps_to_if
def all_zeroes(tree):
if tree == 0:
return True
if not isinstance(tree, Tree):
return False
__tmp0__ = all_zeroes(tree.left)
if __tmp0__:
__tmp0__ = all_zeroes(tree.right)
return __tmp0__
# ------------------------------------
# Pass (4): promote_to_temporary_m
def all_zeroes(tree):
if tree == 0:
return True
if not isinstance(tree, Tree):
return False
__tmp1__ = all_zeroes(tree.left)
__tmp0__ = __tmp1__
if __tmp0__:
__tmp2__ = all_zeroes(tree.right)
__tmp0__ = __tmp2__
return __tmp0__
# ------------------------------------
# Pass (5): remove_trivial_temporaries
def all_zeroes(tree):
if tree == 0:
return True
if not isinstance(tree, Tree):
return False
__tmp0__ = all_zeroes(tree.left)
if __tmp0__:
__tmp0__ = all_zeroes(tree.right)
return __tmp0__
# ------------------------------------
# Pass (6): insert_jumps
def all_zeroes(tree):
if __pc == 0:
if tree == 0:
if __pc == 0:
return True
__pc = 1
if not isinstance(tree, Tree):
if __pc == 0:
return False
__pc = 1
__tmp0__ = all_zeroes(tree.left)
__pc = 1
if __pc == 1:
if __tmp0__:
if __pc == 1:
__tmp0__ = all_zeroes(tree.right)
__pc = 2
__pc = 2
if __pc == 2:
return __tmp0__
__pc = 3
# ------------------------------------
# Pass (7): lift_locals_to_frame
def all_zeroes(tree):
if frame['__pc'] == 0:
if frame['tree'] == 0:
if frame['__pc'] == 0:
return True
frame['__pc'] = 1
if not isinstance(frame['tree'], Tree):
if frame['__pc'] == 0:
return False
frame['__pc'] = 1
frame['__tmp0__'] = all_zeroes(frame['tree'].left)
frame['__pc'] = 1
if frame['__pc'] == 1:
if frame['__tmp0__']:
if frame['__pc'] == 1:
frame['__tmp0__'] = all_zeroes(frame['tree'].right)
frame['__pc'] = 2
frame['__pc'] = 2
if frame['__pc'] == 2:
return frame['__tmp0__']
frame['__pc'] = 3
# ------------------------------------
# Pass (8): add_trampoline_returns
def all_zeroes(tree):
if frame['__pc'] == 0:
if frame['tree'] == 0:
if frame['__pc'] == 0:
frame['__pc'] = 1
return RetOp(value=True)
if not isinstance(frame['tree'], Tree):
if frame['__pc'] == 0:
frame['__pc'] = 1
return RetOp(value=False)
frame['__pc'] = 1
return CallOp(func='all_zeroes', args=[frame['tree'].left], kwargs={}, ret_variable='__tmp0__')
if frame['__pc'] == 1:
if frame['__tmp0__']:
if frame['__pc'] == 1:
frame['__pc'] = 2
return CallOp(func='all_zeroes', args=[frame['tree'].right], kwargs={}, ret_variable='__tmp0__')
frame['__pc'] = 2
if frame['__pc'] == 2:
frame['__pc'] = 3
return RetOp(value=frame['__tmp0__'])
# ------------------------------------
# Pass (9): fix_fn_def
def __fiberfn_all_zeroes(frame):
if frame['__pc'] == 0:
if frame['tree'] == 0:
if frame['__pc'] == 0:
frame['__pc'] = 1
return RetOp(value=True)
if not isinstance(frame['tree'], Tree):
if frame['__pc'] == 0:
frame['__pc'] = 1
return RetOp(value=False)
frame['__pc'] = 1
return CallOp(func='all_zeroes', args=[frame['tree'].left], kwargs={}, ret_variable='__tmp0__')
if frame['__pc'] == 1:
if frame['__tmp0__']:
if frame['__pc'] == 1:
frame['__pc'] = 2
return CallOp(func='all_zeroes', args=[frame['tree'].right], kwargs={}, ret_variable='__tmp0__')
frame['__pc'] = 2
if frame['__pc'] == 2:
frame['__pc'] = 3
return RetOp(value=frame['__tmp0__'])
# ------------------------------------
================================================
FILE: src/expressions.py
================================================
# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import ast
from collections.abc import Container
import utils
import itertools
def promote_call_expressions(expression: ast.AST, fns: Container[str], name_iter, assignments):
"""Given a expression, transforms call expressions to functions in fns by
promoting them to temporary variables. Appends promoting assignment
expressions to assignments. Returns the new expression with call expressions
replaced."""
def map_attributes(field, expression):
return promote_call_expressions(expression, fns, name_iter, assignments)
# Recursively map the expression's attributes (e.g. subexpressions).
new_expression = utils.map_expression(expression, map_attributes)
# If the function is a special function, then lift the expression to a
# temporary. Also, replace the current expression with a reference to that
# temporary.
if isinstance(expression, ast.Call) and isinstance(expression.func, ast.Name) and expression.func.id in fns:
name = next(name_iter)
assignments.append(utils.make_assign(name, new_expression))
return utils.make_lookup(name)
return new_expression
def make_and_if(name: str, expression: ast.AST, body):
return ast.If(
test=utils.make_lookup(name),
body=body,
orelse=[],
)
def make_or_if(name: str, expression: ast.AST, body):
return ast.If(
test=utils.make_not(utils.make_lookup(name)),
body=body,
orelse=[],
)
def promote_boolean_expression_operands(expression: ast.AST, name_iter, lines):
"""Given a expression, transforms boolean expressions by promoting their
operands to temporary values assigned to by if expressions. Returns the
resulting temporary variable, and appends to lines the corresponding if
expressions that populate the variable.
"""
def map_attributes(field, expression):
# If the expression is a boolop, append if statements to lines.
if isinstance(expression, ast.BoolOp):
assert len(expression.values) > 0
maker = make_or_if if isinstance(
expression.op, ast.Or) else make_and_if
name = next(name_iter)
lines.append(utils.make_assign(name, expression.values[0]))
for child in itertools.islice(expression.values, 1, None):
body = []
body.append(promote_boolean_expression_operands(
utils.make_assign(name, child),
name_iter,
body,
))
lines.append(maker(name, child, body))
return utils.make_lookup(name)
return promote_boolean_expression_operands(expression, name_iter, lines)
# Recursively map the expression's attributes (e.g. subexpressions).
return utils.map_expression(expression, map_attributes)
FRAME_LOCAL_NAME = "frame"
def promote_variable_access(expression: ast.AST, name_fn):
def map_attributes(field, expression):
if isinstance(expression, ast.Name) and (name := name_fn(expression.id)) is not None:
return ast.Subscript(
ctx=expression.ctx,
slice=ast.Constant(value=name),
value=ast.Name(id=FRAME_LOCAL_NAME, ctx=ast.Load())
)
return promote_variable_access(expression, name_fn)
return utils.map_expression(expression, map_attributes)
================================================
FILE: src/fiber.py
================================================
# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import ast
from collections import ChainMap
from dataclasses import dataclass
import inspect
from typing import Any, Container, Dict, List, Set, Union
import textwrap
import jumps
import mappers
import utils
def get_tree(fn):
lines = textwrap.dedent(inspect.getsource(fn)).split("\n")
for i, line in enumerate(lines):
if line.startswith("def"):
break
return ast.parse("\n".join(lines[i:]))
def fn_call_names(stmt: ast.AST, fns: Container[str]):
for node in ast.walk(stmt):
if isinstance(node, ast.Call):
if isinstance(node.func, ast.Name) and node.func.id in fns:
yield node.func.id
def make_prev_dict(block: ast.AST):
prev_dict = {}
def helper(block: ast.AST):
assert utils.is_block(block)
for i, stmt in enumerate(block.body):
if 0 <= i - 1:
assert stmt not in prev_dict
prev_dict[stmt] = block.body[i - 1]
for stmt in block.body:
if utils.is_block(stmt):
helper(stmt)
helper(block)
return prev_dict
@dataclass
class CallOp:
func: Any
args: List[Any]
kwargs: Dict[Any, Any]
ret_variable: str
@dataclass
class TailCallOp:
func: Any
args: List[Any]
kwargs: Dict[Any, Any]
@dataclass
class RetOp:
value: Any
OP_MAP = {
CallOp.__name__: CallOp,
TailCallOp.__name__: TailCallOp,
RetOp.__name__: RetOp,
}
def matches_call(call: Union[ast.expr, None], fns: Container[str]):
return isinstance(call, ast.Call) \
and isinstance(call.func, ast.Name) \
and call.func.id in fns
def matches_callop(stmt: ast.AST, fns: Container[str]):
return isinstance(stmt, ast.Assign) and \
len(stmt.targets) == 1 and \
isinstance(target := stmt.targets[0], ast.Subscript) and \
isinstance(target.value, ast.Name) and \
target.value.id == "frame" and \
matches_call(stmt.value, fns)
def matches_tailcallop(stmt: ast.AST, fns: Container[str]):
return isinstance(stmt, ast.Return) and matches_call(stmt.value, fns)
def matches_retop(stmt: ast.AST, fns: Container[str]):
return isinstance(stmt, ast.Return)
def is_pc_assign(stmt: ast.AST):
return isinstance(stmt, ast.Assign) and \
len(stmt.targets) == 1 and \
isinstance(target := stmt.targets[0], ast.Subscript) and \
isinstance(target.value, ast.Name) and \
target.value.id == "frame" and \
isinstance(name := target.slice, ast.Constant) and \
name.value == jumps.PC_LOCAL_NAME and \
isinstance(stmt.value, ast.Constant) and \
isinstance(stmt.value.value, int)
def is_tail_call(stmt: ast.AST):
return isinstance(stmt, ast.Return) and \
isinstance(stmt.value, ast.Call)
def make_callop_expr(variable: ast.Constant, call: ast.Call):
return ast.Return(
value=ast.Call(
func=utils.make_lookup(CallOp.__name__),
args=[],
keywords=[
ast.keyword(arg="func", value=ast.Constant(
value=call.func.id)),
ast.keyword(arg="args", value=ast.List(
elts=call.args, ctx=ast.Load())),
ast.keyword(arg="kwargs", value=ast.Dict(
keys=[ast.Constant(k.arg) for k in call.keywords], values=[k.value for k in call.keywords])),
ast.keyword(arg="ret_variable", value=variable),
]
)
)
def make_tailcallop_expr(call: ast.Call):
return ast.Return(
value=ast.Call(
func=utils.make_lookup(TailCallOp.__name__),
args=[],
keywords=[
ast.keyword(arg="func", value=ast.Constant(
value=call.func.id)),
ast.keyword(arg="args", value=ast.List(
elts=call.args, ctx=ast.Load())),
ast.keyword(arg="kwargs", value=ast.Dict(
keys=[ast.Constant(k.arg) for k in call.keywords], values=[k.value for k in call.keywords])),
]
)
)
def make_retop_expr(value: ast.AST):
return ast.Return(
value=ast.Call(
func=utils.make_lookup(RetOp.__name__),
args=[],
keywords=[
ast.keyword(arg="value", value=value)
]
)
)
def make_arguments():
return ast.arguments(
posonlyargs=[],
args=[ast.arg(arg="frame")],
kwonlyargs=[],
kwarg=None,
vararg=None,
defaults=[],
kw_defaults=[],
)
def fix_fn_def(fn_tree: ast.FunctionDef, fn):
fn_tree.name = f"__fiberfn_{fn.__name__}"
fn_tree.args = make_arguments()
def fiber_locals(fn_tree: ast.FunctionDef):
local_vars = utils.local_vars(fn_tree)
local_vars.add(jumps.PC_LOCAL_NAME)
return local_vars
def insert_jumps(fn_tree: ast.FunctionDef, prev_dict, fns):
prev_dict = make_prev_dict(fn_tree)
body, _ = jumps.insert_jumps(
fn_tree.body, jump_to=lambda stmt: needs_jump(stmt, prev_dict, fns))
return body
def lift_locals_to_frame(fn_tree: ast.FunctionDef):
local_vars = fiber_locals(fn_tree)
return mappers.map_scope(fn_tree, mappers.lift_to_frame_m(
name_fn=lambda x: x if x in local_vars else None))
def needs_jump(stmt: ast.AST, prev_dict, fns):
if not stmt in prev_dict:
return False
# Check whether the child function is a trampoline.
fn_calls = list(fn_call_names(prev_dict[stmt], fns))
if not fn_calls:
return False
for name in fn_calls:
if not (name in FIBER_FN_NAME_MAP or name in fns):
return False
return not is_tail_call(prev_dict[stmt])
def compile_tree(tree: ast.AST, fn, local_vars):
tree = ast.fix_missing_locations(tree)
code = compile(tree, f" {inspect.getfile(fn)}", "exec")
results = {}
exec(code, dict([*fn.__globals__.items(), *
OP_MAP.items(), *local_vars.items()]), results)
results[tree.body[0].name].__fibercode__ = ast.unparse(tree)
return results[tree.body[0].name]
# This is hacky...
@dataclass
class FiberMetadata:
fn_def: ast.FunctionDef
fn: Any
FIBER_FN_NAME_MAP = {}
FIBER_FN_COMPILED_MAP = {}
def add_trampoline_returns(block: ast.AST, fns: Container[str]):
"""Recursively mutates the block by replacing a function call or a return
statement with a return to a trampoline. Also moves the PC assignment to
before the return to the trampoline.
We assume that all recursive calls have been lifted to temporaries, and
tail calls are in `return call()` form (trivial temporary eliminated)."""
assert utils.is_block(block)
# Make a shallow copy, as we mutate the list as we iterate (by swapping).
body = block.body
for index, stmt in enumerate(list(body)):
if utils.is_block(stmt):
add_trampoline_returns(stmt, fns)
continue
if matches_callop(stmt, fns):
replaced = make_callop_expr(stmt.targets[0].slice, stmt.value)
elif matches_tailcallop(stmt, fns):
replaced = make_tailcallop_expr(stmt.value)
elif matches_retop(stmt, fns):
replaced = make_retop_expr(stmt.value)
else:
continue
assert index + 1 < len(body)
assert is_pc_assign(body[index + 1])
body[index], body[index + 1] = body[index+1], replaced
def fiber(fns: Container[str] = None, *, locals, recursive=True):
"""Returns a decorator that converts a function to a fiber.
A fiber is a userspace scheduled thread. In this fiber implementation, we
yield to the userspace scheduler whenever a listed function is called.
Using the trampoline scheduler, we can functions that recurse arbitrarily
deep by simulating the call stack on the heap: Suppose we are executing a
function A. When A reaches a call to some function B in fns, instead of
calling the B directly, A will return a call operation to a trampoline. The
trampoline will call the B with the correct arguments. After B finishes
executing, the trampoline will resume A, passing B's return value.
>>> @fiber(locals=locals())
... def fib(n):
... if n <= 1: return n
... return fib(n-1) + fib(n-2)
...
>>> import trampoline
>>> trampoline.run(fib, [10])
55
"""
if fns is None:
fns = set()
for fiber_fn in FIBER_FN_NAME_MAP:
fns = set(fns)
fns.add(fiber_fn)
if callable(fns):
raise ValueError("Did you forget to call the fiber decorator?")
def make_fiber(fn):
if recursive:
fns.add(fn.__name__)
tree = get_tree(fn)
name_iter, fn_tree = utils.dunder_names(), tree.body[0]
assert isinstance(fn_tree, ast.FunctionDef)
transforms = [
mappers.for_to_while_m(name_iter),
mappers.promote_while_cond_m(name_iter),
mappers.bool_exps_to_if_m(name_iter),
mappers.promote_to_temporary_m(fns, name_iter),
]
for t in transforms:
fn_tree = mappers.map_scope(fn_tree, t)
# These mappers need access to the new tree to preprocess variables.
fn_tree = mappers.map_scope(fn_tree, mappers.remove_trivial_temporaries_m(fn_tree))
prev_dict = make_prev_dict(fn_tree)
fn_tree.body = insert_jumps(fn_tree, prev_dict, fns)
fn_tree = lift_locals_to_frame(fn_tree)
add_trampoline_returns(fn_tree, fns)
fix_fn_def(fn_tree, fn)
tree.body[0] = fn_tree
fiber_fn = compile_tree(tree, fn, locals)
lookup = FiberMetadata(get_tree(fn).body[0], fiber_fn)
FIBER_FN_NAME_MAP[fn.__name__] = lookup
FIBER_FN_COMPILED_MAP[fiber_fn] = lookup
return fiber_fn
return make_fiber
================================================
FILE: src/fiber_test.py
================================================
# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import unittest
import fiber
class TestFiber(unittest.TestCase):
def test_fib(self):
@fiber.fiber(locals=locals())
def fib(n):
if n == 0:
return 0
if n == 1:
return 1
return fib(n-1) + fib(n=n-2)
self.maxDiff = None
want = """
def __fiberfn_fib(frame):
if frame['__pc'] == 0:
if frame['n'] == 0:
if frame['__pc'] == 0:
frame['__pc'] = 1
return RetOp(value=0)
if frame['n'] == 1:
if frame['__pc'] == 0:
frame['__pc'] = 1
return RetOp(value=1)
frame['__pc'] = 1
return CallOp(func='fib', args=[frame['n'] - 1], kwargs={}, ret_variable='__tmp0__')
if frame['__pc'] == 1:
frame['__pc'] = 2
return CallOp(func='fib', args=[], kwargs={'n': frame['n'] - 2}, ret_variable='__tmp1__')
if frame['__pc'] == 2:
frame['__pc'] = 3
return RetOp(value=frame['__tmp0__'] + frame['__tmp1__'])
""".strip()
self.assertEqual(want, fib.__fibercode__)
def test_sum(self):
@fiber.fiber(locals=locals())
def sum(lst, acc):
if not lst:
return acc
return sum(lst[1:], acc + lst[0])
want = """
def __fiberfn_sum(frame):
if frame['__pc'] == 0:
if not frame['lst']:
if frame['__pc'] == 0:
frame['__pc'] = 1
return RetOp(value=frame['acc'])
frame['__pc'] = 1
return TailCallOp(func='sum', args=[frame['lst'][1:], frame['acc'] + frame['lst'][0]], kwargs={})
""".strip()
self.assertEqual(want, sum.__fibercode__)
if __name__ == '__main__':
unittest.main()
================================================
FILE: src/jumps.py
================================================
# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import ast
from typing import List, Callable, Iterable
import utils
PC_LOCAL_NAME = "__pc"
def is_supported_jump_block(b: ast.AST):
return any(isinstance(b, t) for t in (ast.While, ast.If))
def block_has_jump_to(block, jump_to: Callable[[ast.AST], bool]):
return any(has_jump_to(stmt, jump_to) for stmt in block.body)
def has_jump_to(stmt: ast.AST, jump_to: Callable[[ast.AST], bool]):
return jump_to(stmt) or (is_supported_jump_block(stmt) and block_has_jump_to(stmt, jump_to))
def partition_stmts(stmts: Iterable[ast.AST], jump_to: Callable[[ast.AST], bool]):
"""Splits statements into partitions with jump_to statements as dividers.
The first statement in each partition is one that needs to be jumped to."""
current = []
for stmt in stmts:
if has_jump_to(stmt, jump_to):
yield current
current = []
current.append(stmt)
yield current
def transform_if(stmt: ast.If, jump_to: Callable[[ast.AST], bool], next_pc):
body, next_pc = insert_jumps(
stmt.body, jump_to, next_pc)
return ast.If(test=stmt.test, body=body, orelse=stmt.orelse), next_pc
def transform_while(stmt: ast.While, jump_to: Callable[[ast.AST], bool], next_pc):
first_pc = next_pc
body, next_pc = insert_jumps(stmt.body, jump_to, next_pc)
# While loops jump back to the start of the loop.
body.append(utils.make_assign(PC_LOCAL_NAME, ast.Constant(first_pc)))
return ast.While(test=stmt.test, body=body, orelse=stmt.orelse), next_pc
def make_range_test(start_pc, end_pc):
"""Creates an boolean expression AST that checks whether the pc variable is
in range(start, end)."""
if start_pc + 1 == end_pc:
return ast.Compare(
left=utils.make_lookup(PC_LOCAL_NAME),
ops=[ast.Eq()],
comparators=[ast.Constant(value=start_pc)]
)
return ast.Compare(
left=ast.Constant(value=start_pc),
ops=[ast.LtE(), ast.Lt()],
comparators=[utils.make_lookup(
PC_LOCAL_NAME), ast.Constant(value=end_pc)],
)
def transform_partition(partition, jump_to, next_pc):
"""Recursively transforms the partition, and wraps it in the appropriate if
statement. Returns the new AST as well as the next pc value."""
body = []
start_pc = next_pc
next_pc += 1 # Need at least one PC for this partition.
for stmt in partition:
transformed = stmt
if isinstance(stmt, ast.If):
# For blocks, the first inner PC is the same PC as the outer PC.
transformed, next_pc = transform_if(stmt, jump_to, next_pc-1)
elif isinstance(stmt, ast.While):
transformed, next_pc = transform_while(stmt, jump_to, next_pc-1)
body.append(transformed)
end_pc = next_pc
body.append(utils.make_assign(PC_LOCAL_NAME, ast.Constant(end_pc)))
return ast.If(test=make_range_test(start_pc, end_pc), body=body, orelse=[]), next_pc
def insert_jumps(stmts: Iterable[ast.AST], jump_to: Callable[[ast.AST], bool], start_pc=0):
"""Inserts ifs into a sequence of statements such that each statement for
which jump_to returns True has a pc value where entering the sequence with
that value jumps to that statement.
Only recursively inserts jumps inside child if and for blocks. If your
block has for loops, then use for_to_while_m to rewrite them to while.
Returns a new list of statements and the total number of jumps inserted.
"""
new_stmts = []
next_pc = start_pc
for partition in partition_stmts(stmts, jump_to):
if not partition:
continue
transformed, next_pc = transform_partition(partition, jump_to, next_pc)
new_stmts.append(transformed)
return new_stmts, next_pc
================================================
FILE: src/jumps_test.py
================================================
# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import ast
import unittest
import jumps
class TestInsertJumps(unittest.TestCase):
def test_complex(self):
source = """
def bar(hello, world):
for i in range(10):
print(hello)
a = hello + world
while b < a:
print()
b = a * a
a = a + b
if True:
d = a + b
d = d * 2
d = d + 1
d = d + 2
return d
""".strip()
want = """
def bar(hello, world):
if __pc == 0:
for i in range(10):
print(hello)
__pc = 1
if __pc == 1:
a = hello + world
__pc = 2
if 2 <= __pc < 5:
while b < a:
if __pc == 2:
print()
__pc = 3
if __pc == 3:
b = a * a
__pc = 4
if __pc == 4:
a = a + b
__pc = 5
__pc = 2
__pc = 5
if 5 <= __pc < 7:
if True:
if __pc == 5:
d = a + b
__pc = 6
if __pc == 6:
d = d * 2
__pc = 7
__pc = 7
if __pc == 7:
d = d + 1
__pc = 8
if __pc == 8:
d = d + 2
return d
__pc = 9
""".strip()
tree = ast.parse(source)
fn_tree = tree.body[0]
assert isinstance(fn_tree, ast.FunctionDef)
fn_tree.body, _ = jumps.insert_jumps(
fn_tree.body, lambda stmt: isinstance(stmt, ast.Assign))
tree = ast.fix_missing_locations(tree)
result = ast.unparse(tree)
self.assertEqual(result, want)
def test_equivalent(self):
source = """
def fib(__pc, n):
curr = 0
next = 1
for i in range(n):
tmp = curr
curr = next
next = tmp + next
return curr
"""
tree = ast.parse(source)
fn_tree = tree.body[0]
assert isinstance(fn_tree, ast.FunctionDef)
fn_tree.body, _ = jumps.insert_jumps(
fn_tree.body, lambda stmt: isinstance(stmt, ast.Assign))
fn_tree.name = "fib_transformed"
tree = ast.fix_missing_locations(tree)
code = compile(tree, "", "exec")
results = {}
exec(code, globals(), results)
exec(source, globals(), results)
fib, fib_transformed = results["fib"], results["fib_transformed"]
for i in range(100):
self.assertEqual(fib(0, i), fib_transformed(0, i))
if __name__ == '__main__':
unittest.main()
================================================
FILE: src/manual.py
================================================
# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from dataclasses import dataclass
from typing import Any
def editDistance(first, second):
"""Computes the edit distance between two strings.
>>> editDistance("kitten", "sitting")
3
"""
return editDistanceImpl(first, second, 0, 0)
def editDistanceImpl(first, second, f, s):
if f == len(first):
return len(second) - s
if s == len(second):
return len(first) - f
if first[f] == second[s]:
return editDistanceImpl(first, second, f+1, s+1)
deleteFirst = editDistanceImpl(first, second, f+1, s) + 1
deleteSecond = editDistanceImpl(first, second, f, s+1) + 1
replace = editDistanceImpl(first, second, f+1, s+1) + 1
return min(deleteFirst, deleteSecond, replace)
@dataclass
class Frame:
__pc: int
locals: Any
child: Any
@dataclass
class CallOp:
fn: Any
arguments: Any
@dataclass
class TailCallOp:
callop: CallOp
@dataclass
class RetOp:
pass
def engine(fn, args):
stack = [] # make frame fn, args
while True:
# op = run last function on stack
# if op == tailcall
# pop the last frame in stack
# if op == call or tailcall
# bind arguments to corresponding frame locals
# push new frame onto stack
# if op == ret
# pop last frame in stack
# if no parent frame, return value
# assign new top's child frame pointer to popped frame
# next function will read the return value through the pointer
# ABI: returns put return value into the current frame
# We don't support nested functions.
def editDistanceIterImpl(frame):
if frame.__pc == 0:
if frame.f == len(frame.first):
return len(frame.second) - frame.s
if frame.s == len(frame.second):
return len(frame.first) - frame.s
if 0 <= frame.__pc <= 1:
if frame.first[frame.f] == frame.second[frame.s]:
if frame.__pc == 0:
frame.__pc = 1
# callop editDistanceImpl(frame.first, frame.second, frame.f+1, frame.s+1)
# jmp (1)
return frame.child.ret
frame.__pc = 2
if frame.__pc == 2:
frame.__pc = 3
# callop editDistanceImpl(frame.first, frame.second, frame.f+1, frame.s)
frame.deleteFirst = frame.child.ret + 1
# jmp
if frame.__pc == 3:
frame.__pc = 4
# callop editDistanceImpl(frame.first, frame.second, frame.f, frame.s+1)
frame.deleteSecond = frame.child.ret + 1
# jmp
if frame.__pc == 4:
frame.__pc = 5
# callop editDistanceImpl(frame.first, frame.second, frame.f+1, frame.s+1)
frame.replace = frame.child.ret + 1
# jmp
if frame.__pc == 5:
return min(frame.deleteFirst, frame.deleteSecond, frame.replace)
# transforms
# mark recursive calls
# promote nested recursive call results to temporaries
# change locals to accesses in heap frame
# add jump points after recursive calls
# write recursive calls as returning (tail)?call ops
# write returns as returning return ops
# write jumps as if statements & pc counter
================================================
FILE: src/mappers.py
================================================
# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import ast
from collections.abc import Container
import expressions
import utils
def map_scope(scope, fn):
"""Applies the mapping function to every statement in the scope.
The mapping function should return a list of statements; the returned
statements will be flattened together."""
kwargs = {field: value for field, value in ast.iter_fields(scope)}
body = []
for stmt in scope.body:
body.extend(fn(stmt))
# TODO(tylerhou): Should we map all scopes?
if utils.is_supported_scope(stmt):
body[-1] = map_scope(body[-1], fn)
kwargs["body"] = body
return type(scope)(**kwargs)
def promote_to_temporary_m(fns: Container[str], name_iter):
"""Creates a function mapper that promotes the results of inner calls to
functions in fns to temporary variables."""
def promote_mapper(stmt):
stmts = []
stmts.append(expressions.promote_call_expressions(
stmt, fns, name_iter, stmts)) # Mutates stmts
return stmts
return promote_mapper
def remove_trivial_temporaries_m(fn_ast: ast.AST):
"""Creates a function mapper that removes trivial assignments."""
trivial_temps = set(utils.potentially_trivial_temporaries(fn_ast))
first_assignment, last_assignment = utils.find_assignments(fn_ast, trivial_temps)
trivial_temps = set(t for t in trivial_temps if last_assignment[t] is first_assignment[t])
trivial_assignments = set(last_assignment[t] for t in trivial_temps)
to_replace = {temp: last_assignment[temp] for temp in trivial_temps if temp in trivial_temps}
def remove_trivial_mapper(stmt):
return [] if stmt in trivial_assignments else [utils.replace_variable(stmt, to_replace)]
return remove_trivial_mapper
def for_to_while_m(name_iter):
"""Creates a function mapper that converts for loops to equivalent while loops."""
def mapper(stmt):
if not isinstance(stmt, ast.For):
return [stmt]
iter_n, test_n = next(name_iter), next(name_iter)
body = [utils.make_for_try(stmt.target, iter_n, test_n)] + stmt.body
return [
utils.make_assign(iter_n, utils.make_call("iter", stmt.iter)),
utils.make_assign(test_n, ast.Constant(value=True)),
ast.While(test=utils.make_lookup(test_n),
body=body, orelse=stmt.orelse),
]
return mapper
def promote_while_cond_m(name_iter):
"""Creates a function mapper that promotes the test in while loops to a variable."""
def mapper(stmt):
if not isinstance(stmt, ast.While):
return [stmt]
if isinstance(stmt.test, ast.Name):
return [stmt]
# TODO(tylerhou): Add a test for this.
condition_n = next(name_iter)
test_assign = utils.make_assign(condition_n, stmt.test)
body = stmt.body + [test_assign]
return [test_assign, ast.While(test=utils.make_lookup(condition_n), body=body, orelse=stmt.orelse)]
return mapper
def bool_exps_to_if_m(name_iter):
"""Creates a function mapper that rewrites boolean expressions as if
statements so promotion to temporaries doesn't change evaluation order."""
def mapper(stmt):
stmts = []
stmts.append(expressions.promote_boolean_expression_operands(
stmt, name_iter, stmts)) # Mutates stmts
return stmts
return mapper
def lift_to_frame_m(name_fn=lambda x: x):
"""Creates a function mapper that replaces all accesses where name_fn
returns not None to loads and stores in a frame object."""
def mapper(stmt):
return [expressions.promote_variable_access(stmt, name_fn)]
return mapper
================================================
FILE: src/mappers_test.py
================================================
# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import ast
import unittest
import mappers
import utils
def map_function(source, mapper):
tree = ast.parse(source).body[0]
tree = mappers.map_scope(tree, mapper)
tree = ast.fix_missing_locations(tree)
return ast.unparse(tree)
class TestMappers(unittest.TestCase):
def test_promote_to_temporary_m(self):
source = """
def foo():
t = bar(2)
if t == 1:
return bar(bar(bar(1), 3), 5) + bar(baz(t), 4)
if t == bar(10):
return bar(baz(t), 2)
""".strip()
want = """
def foo():
__tmp0__ = bar(2)
t = __tmp0__
if t == 1:
__tmp1__ = bar(1)
__tmp2__ = bar(__tmp1__, 3)
__tmp3__ = bar(__tmp2__, 5)
__tmp4__ = bar(baz(t), 4)
return __tmp3__ + __tmp4__
__tmp5__ = bar(10)
if t == __tmp5__:
__tmp6__ = bar(baz(t), 2)
return __tmp6__
""".strip()
mapper = mappers.promote_to_temporary_m(["bar"], utils.dunder_names())
result = map_function(source, mapper)
self.assertEqual(result, want)
def test_remove_trivial_temporaries_m(self):
source = """
def foo():
__tmp0__ = bar(2)
t = __tmp0__
if t == 1:
__tmp1__ = bar(1)
__tmp2__ = bar(__tmp1__, 3)
__tmp3__ = bar(__tmp2__, 5)
__tmp4__ = bar(baz(t), 4)
return __tmp3__ + __tmp4__
__tmp5__ = bar(10)
if t == __tmp5__:
__tmp6__ = bar(baz(t), 2)
return __tmp6__
""".strip()
want = """
def foo():
t = bar(2)
if t == 1:
__tmp1__ = bar(1)
__tmp2__ = bar(__tmp1__, 3)
__tmp3__ = bar(__tmp2__, 5)
__tmp4__ = bar(baz(t), 4)
return __tmp3__ + __tmp4__
__tmp5__ = bar(10)
if t == __tmp5__:
return bar(baz(t), 2)
""".strip()
tree = ast.parse(source).body[0]
# Remove trivial needs to preprocess the tree to find trivial variables.
mapper = mappers.remove_trivial_temporaries_m(tree)
tree = mappers.map_scope(tree, mapper)
tree = ast.fix_missing_locations(tree)
result = ast.unparse(tree)
self.assertEqual(result, want)
def test_remove_trivial_temporaries_m_tail_call(self):
source = """
def sum(lst, acc):
if not lst:
return acc
__tmp0__ = sum(lst[1:], acc + lst[0])
return __tmp0__
""".strip()
want = """
def sum(lst, acc):
if not lst:
return acc
return sum(lst[1:], acc + lst[0])
""".strip()
tree = ast.parse(source).body[0]
mapper = mappers.remove_trivial_temporaries_m(tree)
tree = mappers.map_scope(tree, mapper)
tree = ast.fix_missing_locations(tree)
result = ast.unparse(tree)
self.assertEqual(result, want)
def test_for_to_while_m(self):
source = """
def bar():
pre = 1
for i in range(10):
print(pre, i)
if i == 5:
break
else:
print('else')
post = 1
return pre + post
""".strip()
want = """
def bar():
pre = 1
__tmp0__ = iter(range(10))
__tmp1__ = True
while __tmp1__:
try:
i = next(__tmp0__)
except StopIteration:
__tmp1__ = False
continue
print(pre, i)
if i == 5:
break
else:
print('else')
post = 1
return pre + post
""".strip()
mapper = mappers.for_to_while_m(utils.dunder_names())
result = map_function(source, mapper)
self.assertEqual(result, want)
def test_promote_while_cond_m(self):
source = """
def bar():
p = [1, 2, 3]
while len(p) > 0:
t = p.pop()
else:
print('else')
post = 1
return t + post
""".strip()
want = """
def bar():
p = [1, 2, 3]
__tmp0__ = len(p) > 0
while __tmp0__:
t = p.pop()
__tmp0__ = len(p) > 0
else:
print('else')
post = 1
return t + post
""".strip()
mapper = mappers.promote_while_cond_m(utils.dunder_names())
result = map_function(source, mapper)
self.assertEqual(result, want)
def test_bool_exps_to_if_m(self):
source = """
def bar():
a = first() and (second() or third()) and fourth()
b = 1 + (foo() or baz())
return a or b
""".strip()
want = """
def bar():
__tmp0__ = first()
if __tmp0__:
__tmp1__ = second()
if not __tmp1__:
__tmp1__ = third()
__tmp0__ = __tmp1__
if __tmp0__:
__tmp0__ = fourth()
a = __tmp0__
__tmp2__ = foo()
if not __tmp2__:
__tmp2__ = baz()
b = 1 + __tmp2__
__tmp3__ = a
if not __tmp3__:
__tmp3__ = b
return __tmp3__
""".strip()
mapper = mappers.bool_exps_to_if_m(utils.dunder_names())
result = map_function(source, mapper)
self.assertEqual(result, want)
def test_lift_to_frame_m(self):
source = """
def bar(arg1, arg2):
a = arg1 + arg2
if a == 2:
return arg1 + arg2
b = a + (foo() or baz())
return a + b
""".strip()
want = """
def bar(arg1, arg2):
frame['a'] = frame['arg1'] + frame['arg2']
if frame['a'] == 2:
return frame['arg1'] + frame['arg2']
frame['b'] = frame['a'] + (foo() or baz())
return frame['a'] + frame['b']
""".strip()
mapper = mappers.lift_to_frame_m(
lambda x: x if x in ("a", "b", "arg1", "arg2") else None)
result = map_function(source, mapper)
self.assertEqual(result, want)
if __name__ == '__main__':
unittest.main()
================================================
FILE: src/trampoline.py
================================================
# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import ast
from dataclasses import dataclass
from typing import Any, Dict, List, Union
import itertools
import fiber
import jumps
@dataclass
class StackFrame:
frame: Dict[str, Any]
fn: Any
ret_variable: Union[str, None]
def pos_with_defaults(args: ast.arguments):
num_non_defaults = len(args.posonlyargs) + \
len(args.args) - len(args.defaults)
args_iter = itertools.chain(iter(args.posonlyargs), iter(args.args))
for _ in range(num_non_defaults):
yield next(args_iter), None
for default in args.defaults:
yield next(args_iter), default
def bind_frame(positional_args, keyword_args, fn_tree: ast.FunctionDef):
frame = {}
# Reverse the list of args because we pop args from the front of the
# original list, which is the back of the reversed list.
positional_args, fn_args = list(reversed(positional_args)), fn_tree.args
for needed_arg, default in pos_with_defaults(fn_args):
name: str = needed_arg.arg
if positional_args: # If we can still take args
frame[name] = positional_args.pop()
continue
if name in keyword_args:
if needed_arg in fn_tree.args.posonlyargs:
raise TypeError(
f"{fn_tree.name}: got positional only argument {name} as a keyword")
frame[name] = keyword_args[name]
del keyword_args[name]
continue
elif not default:
raise TypeError(f"{fn_tree.name} had too few positional arguments")
frame[name] = default
if fn_args.vararg:
# Unreverse the list as from the beginning.
frame[fn_args.vararg.arg] = list(reversed(positional_args))
keyword_arg_names = set(k.arg for k in itertools.chain(
fn_args.args, fn_args.kwonlyargs))
for kwarg in keyword_args:
if kwarg not in keyword_arg_names:
raise TypeError(
f"{fn_tree.name} got invalid keyword argument '{kwarg}'")
if kwarg in frame:
raise TypeError(
f"{fn_tree.name} got multiple values for argument '{kwarg}'")
frame[kwarg] = keyword_args[kwarg]
for kwarg, default in zip(fn_args.kwonlyargs, fn_args.kw_defaults):
if default is None and kwarg.arg not in frame:
raise TypeError(
f"{fn_tree.name} missing required keyword only argument '{kwarg.arg}'")
if default is not None and kwarg.arg not in frame:
frame[kwarg.arg] = ast.literal_eval(default)
frame[jumps.PC_LOCAL_NAME] = 0
return frame
def run(fn, args=None, kwargs=None, *, __max_stack_size=float('inf')):
if args is None:
args = []
if kwargs is None:
kwargs = {}
frame = bind_frame(args, kwargs, fiber.FIBER_FN_COMPILED_MAP[fn].fn_def)
stack: List[StackFrame] = [StackFrame(frame, fn, None)]
while True:
assert len(stack) <= __max_stack_size
top = stack[-1]
op = top.fn(top.frame)
if isinstance(op, fiber.CallOp):
metadata = fiber.FIBER_FN_NAME_MAP[op.func]
top.ret_variable = op.ret_variable
frame = bind_frame(op.args, op.kwargs, metadata.fn_def)
stack.append(StackFrame(frame, metadata.fn, None))
elif isinstance(op, fiber.TailCallOp):
stack.pop() # Tail call, so we can discard the frame.
metadata = fiber.FIBER_FN_NAME_MAP[op.func]
frame = bind_frame(op.args, op.kwargs, metadata.fn_def)
stack.append(StackFrame(frame, metadata.fn, None))
elif isinstance(op, fiber.RetOp):
stack.pop()
if not stack:
return op.value
top = stack[-1]
assert top.ret_variable is not None
top.frame[top.ret_variable] = op.value
================================================
FILE: src/trampoline_test.py
================================================
# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import sys
import unittest
import fiber
import trampoline
class TestTrampoline(unittest.TestCase):
def test_fib(self):
cache = {}
@fiber.fiber(locals=locals())
def fib(n):
if n in cache:
return cache[n]
if n == 0:
return 0
if n == 1:
return 1
cache[n] = fib(n-1) + fib(n=n-2)
return cache[n]
self.assertEqual(55, trampoline.run(fib, [10], {}))
self.assertLess(0, trampoline.run(fib, [1002], {}))
def test_sum(self):
@fiber.fiber(locals=locals())
def sum(lst, acc):
if not lst:
return acc
return sum(lst[1:], acc + lst[0])
n = sys.getrecursionlimit() + 1
want = n * (n + 1) / 2
got = trampoline.run(sum, [list(range(1, n+1)), 0], __max_stack_size=1)
self.assertEqual(want, got)
def test_sum_recursion_exceeded(self):
def sum(lst, acc):
if not lst:
return acc
return sum(lst[1:], acc + lst[0])
n = sys.getrecursionlimit() + 1
self.assertRaises(RecursionError, sum, list(range(1, n+1)), 0)
def test_sum_non_tailcall(self):
@fiber.fiber(locals=locals())
def sum(lst, acc):
if not lst:
return acc
return sum(lst[1:], acc + lst[0]) + 1
n = sys.getrecursionlimit() + 1
want = n * (n + 1) / 2 + n
got = trampoline.run(sum, [list(range(1, n+1)), 0])
self.assertEqual(want, got)
def test_mutual_recursion(self):
@fiber.fiber(["b"], locals=locals())
def a(n):
if n == 0:
return 1
return b(n-1) * 2
@fiber.fiber(locals=locals())
def b(n):
if n == 0:
return 1
return a(n-1) * 3
got = trampoline.run(a, [10])
self.assertEqual(2**5 * 3**5, got)
def test_edit_distance(self):
def edit_distance(first, second):
@fiber.fiber(locals=locals())
def edit_distance__helper(f, s):
if f == len(first):
return len(second) - s
if s == len(second):
return len(first) - f
if first[f] == second[s]:
return edit_distance__helper(f+1, s+1)
del_f = edit_distance__helper(f+1, s) + 1
replace = edit_distance__helper(f+1, s+1) + 1
del_s = edit_distance__helper(f, s+1) + 1
return min(del_f, replace, del_s)
return trampoline.run(edit_distance__helper, [0, 0])
self.assertEqual(3, edit_distance("kitten", "sitting"))
def test_pop_balloons(self):
def pop_balloons(balloons):
balloons = [1] + balloons + [1]
@fiber.fiber(locals=locals())
def pop_balloons__helper(i, j):
if i == j:
return 0
max_profit = float('-inf')
for k in range(i, j):
profit = pop_balloons__helper(i, k)
profit += pop_balloons__helper(k+1, j)
profit += balloons[i-1] * balloons[k] * balloons[j]
max_profit = max(max_profit, profit)
return max_profit
return trampoline.run(pop_balloons__helper, [1, len(balloons) - 1])
self.assertEqual(175, pop_balloons([4, 5, 7]))
def test_tree_recursion(self):
from collections import namedtuple
Tree = namedtuple("Tree", ["left", "right"])
@fiber.fiber(locals=locals())
def all_zeroes(tree):
if tree == 0:
return True
if not isinstance(tree, Tree):
return False
return all_zeroes(tree.left) and all_zeroes(tree.right)
zeroes = Tree(Tree(0, 0), Tree(Tree(Tree(0, 0), 0), Tree(0, 0)))
one = Tree(Tree(0, 0), Tree(Tree(Tree(1, 0), 0), Tree(0, 0)))
self.assertTrue(trampoline.run(all_zeroes, [zeroes]))
self.assertFalse(trampoline.run(all_zeroes, [one]))
if __name__ == '__main__':
unittest.main()
================================================
FILE: src/utils.py
================================================
# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import ast
import itertools
import re
# We don't support exceptions (ast.Try) or async for/while.
# Exception support is fairly easy to add; async is harder.
FN_INNER_SCOPE_NODES = [ast.For, ast.While, ast.If, ast.While]
def dunder_names():
for i in itertools.count():
yield f"__tmp{i}__"
dunder_regex = re.compile(r'__tmp\d+__')
def is_temporary(name):
return re.match(dunder_regex, name) != None
def is_supported_scope(tree):
return any(isinstance(tree, scope_type) for scope_type in FN_INNER_SCOPE_NODES)
def map_expression(statement: ast.AST, fn):
"""Recursively transforms an expression by applying the mapping function to
all attributes that are AST nodes.
Does not recursively transform scope bodies (while, for, try) as one would
usually call this function from a mapper in map_scope, which itself already
iterates through scope bodies."""
kwargs = {}
for field, value in ast.iter_fields(statement):
result = value
if is_supported_scope(statement) and field == "body":
result = value
elif isinstance(value, list):
result = [fn(field, v)
for v in value if isinstance(v, ast.AST)]
elif isinstance(value, ast.AST):
result = fn(field, value)
kwargs[field] = result
return type(statement)(**kwargs)
def iter_scope(scope):
"""Iterates through every statement in the scope, recursively."""
for stmt in scope.body:
yield stmt
if is_supported_scope(stmt):
yield from iter_scope(stmt)
def potentially_trivial_temporaries(fn):
"""Yields potentially trivial temporaries.
Temporaries are trivial if they are directly assigned to another variable
or if they are returned; e.g.
t = __tmp1__
return __tmp2__
If they are assigned to multiple times, then they are not trivial.
"""
for statement in iter_scope(fn):
if isinstance(statement, ast.Return) and \
isinstance(statement.value, ast.Name) and \
is_temporary(statement.value.id):
yield statement.value.id
if (isinstance(statement, ast.Assign) or
isinstance(statement, ast.AnnAssign)) and \
isinstance(statement.value, ast.Name) and \
is_temporary(statement.value.id):
yield statement.value.id
def find_assignments(fn, variables):
"""Finds the first and last assignment for each variable."""
first_assignment, last_assignment = {}, {}
for statement in iter_scope(fn):
if isinstance(statement, ast.Assign) and \
len(statement.targets) == 1 and \
isinstance(statement.targets[0], ast.Name) and \
(var := statement.targets[0].id) in variables:
if var not in first_assignment:
first_assignment[var] = statement
last_assignment[var] = statement
return first_assignment, last_assignment
def replace_variable(statement, assignments):
def mapper(field, expression):
if field == "value":
return assignments[expression.id].value
return expression
if isinstance(statement, ast.Return) and \
isinstance(statement.value, ast.Name) and \
statement.value.id in assignments:
return map_expression(statement, mapper)
if (isinstance(statement, ast.Assign) or
isinstance(statement, ast.AnnAssign)) and \
isinstance(statement.value, ast.Name) and \
statement.value.id in assignments:
return map_expression(statement, mapper)
return statement
def make_assign(target, value):
return ast.Assign(targets=[ast.Name(id=target, ctx=ast.Store())], value=value)
def make_call(func: str, *args):
return ast.Call(func=ast.Name(id=func, ctx=ast.Load()), args=list(args), keywords=[])
def make_lookup(name: str):
return ast.Name(id=name, ctx=ast.Load())
def make_not(exp: ast.AST):
return ast.UnaryOp(op=ast.Not(), operand=exp)
def make_for_try(loop_target, iter_n, test_n):
return ast.Try(
body=[ast.Assign(targets=[loop_target], value=make_call(
"next", make_lookup(iter_n)))],
handlers=[
ast.ExceptHandler(type=make_lookup("StopIteration"),
body=[
make_assign(
test_n, ast.Constant(value=False)),
ast.Continue()
])],
orelse=[],
finalbody=[])
def is_block(block: ast.AST):
return hasattr(block, "body") and isinstance(block.body, list)
def _locals_impl(fn: ast.AST):
assert isinstance(fn, ast.FunctionDef)
args = fn.args
for arg in itertools.chain(args.posonlyargs, args.args, args.kwonlyargs, (args.vararg, args.kwarg)):
if arg is not None:
yield arg.arg
def helper(block: ast.AST):
for node in block.body:
if isinstance(node, ast.Assign):
for target in node.targets:
if isinstance(target, ast.Name):
yield target.id
if isinstance(node, ast.AnnAssign) or \
isinstance(node, ast.AugAssign) or \
isinstance(node, ast.NamedExpr):
if isinstance(node.target, ast.Name):
yield node.target.id
if is_block(node) and not isinstance(node, ast.FunctionDef):
yield from helper(node)
yield from helper(fn)
def local_vars(fn: ast.AST):
"""Returns a set of all function local variables."""
return set(_locals_impl(fn))