Add register and flag setting logic