Clean up In-Memory platform

Changes

* separate code into different files
* remove TODOs
* clean up functions that are messy