Refactor map implementation