Move some code around to fix a few issues with #15, #16, #18