Issue 504 partial fix.

Looks like I was able to fix second test case from #504. Two other test case are still not produce correct JS code.

I'm not very sure with this update. Please check it carefully before merge. Looks like it fixes problem and doesn't break any existing test.