Fix regression from issue 3312

This fixes a regression caused by the re-factoring in #3312. This PR thus fixes #3382 and #3383.

@yurydelendik Is this fix OK? Or should more of the code be rewritten/re-factored to solve this issue?
If this is OK, I'll attach a test file to this PR as well.