diff --git a/src/tracerestrict.cpp b/src/tracerestrict.cpp index 2e5cc5b0ac..de56f64687 100644 --- a/src/tracerestrict.cpp +++ b/src/tracerestrict.cpp @@ -844,6 +844,16 @@ static uint32 GetDualInstructionInitialValue(TraceRestrictItem item) } } +template T InstructionIteratorNext(T iter) +{ + return IsTraceRestrictDoubleItem(*iter) ? iter + 2 : iter + 1; +} + +template void InstructionIteratorAdvance(T &iter) +{ + iter = InstructionIteratorNext(iter); +} + CommandCost TraceRestrictProgramRemoveItemAt(std::vector &items, uint32 offset, bool shallow_mode) { TraceRestrictItem old_item = *TraceRestrictProgram::InstructionAt(items, offset); @@ -861,10 +871,10 @@ CommandCost TraceRestrictProgramRemoveItemAt(std::vector &ite uint32 recursion_depth = 1; std::vector::iterator remove_start = TraceRestrictProgram::InstructionAt(items, offset); - std::vector::iterator remove_end = remove_start + 1; + std::vector::iterator remove_end = InstructionIteratorNext(remove_start); // iterate until matching end block found - for (; remove_end != items.end(); ++remove_end) { + for (; remove_end != items.end(); InstructionIteratorAdvance(remove_end)) { TraceRestrictItem current_item = *remove_end; if (IsTraceRestrictConditional(current_item)) { if (GetTraceRestrictCondFlags(current_item) == 0) { @@ -873,8 +883,13 @@ CommandCost TraceRestrictProgramRemoveItemAt(std::vector &ite recursion_depth--; if (recursion_depth == 0) { if (remove_whole_block) { - // inclusively remove up to here - ++remove_end; + if (shallow_mode) { + // must erase endif first, as it is later in the vector + items.erase(remove_end, InstructionIteratorNext(remove_end)); + } else { + // inclusively remove up to here + InstructionIteratorAdvance(remove_end); + } break; } else { // exclusively remove up to here @@ -897,27 +912,18 @@ CommandCost TraceRestrictProgramRemoveItemAt(std::vector &ite return_cmd_error(STR_TRACE_RESTRICT_ERROR_CAN_T_SHALLOW_REMOVE_IF_ELIF); } } - } else if (IsTraceRestrictDoubleItem(current_item)) { - // this is a double-item, jump over the next item as well - ++remove_end; } } if (recursion_depth != 0) return CMD_ERROR; // ran off the end if (shallow_mode) { - // must erase endif first, as it is later in the vector - if (remove_whole_block) items.erase(remove_end - 1); - items.erase(remove_start); + items.erase(remove_start, InstructionIteratorNext(remove_start)); } else { items.erase(remove_start, remove_end); } } else { std::vector::iterator remove_start = TraceRestrictProgram::InstructionAt(items, offset); - std::vector::iterator remove_end = remove_start + 1; + std::vector::iterator remove_end = InstructionIteratorNext(remove_start); - if (IsTraceRestrictDoubleItem(old_item)) { - // this is a double-item, remove the next item as well - ++remove_end; - } items.erase(remove_start, remove_end); } return CommandCost();