Memory management refactor