Extract step_index helper and use it in all inner-loop step sites

Extract step_index helper and use it in all inner-loop step sites #714