More changes to staging logic