Instead of checking for stage1 at every callsite, move the logic
inside `allocateAtom`. This is fine since this logic will disappear
anyhow once I add expanding and shifting segments and sections.