Updated memory management code