[BranchFolding] Set correct mem refs (2nd try)