Fix context structure to match store/restore code