[BranchFolding] Set correct mem refs