It's still missing ast optimization and some code cleanup is also needed, but it should be working in the general case. Also write a small unit test for it.
