This patch merges the changes from #2344 and #2491 and do some cleanup In particular, this visit the init code expression when visiting the root element expression Closes #2344 Closes #2491 Fixes #2487