I created a PR with a fix, it would be awesome if you can provide a more complex test case in order to test it.
Thanks