This also fixes the warning about an unused boolean variable.
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.