Better fix for #58, passes tests