[docs]classSwitchExpressionSimplifier(SequenceWalker):""" Identifies switch expressions that adds or minuses a constant, removes the constant from the switch expression, and adjust all case expressions accordingly. """
def_handle_SwitchCase(self,node:SwitchCaseNode,**kwargs):changed=Falseswitch_expr=node.switch_exprconvert=None# pylint: disable=import-outside-toplevelfrom..peephole_optimizations.remove_noop_conversionsimportRemoveNoopConversionswhileisinstance(switch_expr,ailment.Expr.Convert):optimized=RemoveNoopConversions(None,None).optimize(switch_expr)ifoptimizedisnotNone:switch_expr=optimizedcontinueconvert=switch_exprswitch_expr=switch_expr.operandifisinstance(switch_expr,ailment.Expr.BinaryOp):ifswitch_expr.opin{"Add","Sub"}andisinstance(switch_expr.operands[1],ailment.Expr.Const):v=switch_expr.operands[1].valueifswitch_expr.op=="Add":v=-vnew_switch_expr=switch_expr.operands[0]ifconvertisnotNone:# unpack if necessaryifisinstance(new_switch_expr,ailment.Expr.Convert):ifnew_switch_expr.to_bits==convert.from_bits:new_switch_expr=new_switch_expr.operandelse:new_switch_expr=ailment.Expr.Convert(None,new_switch_expr.from_bits,convert.to_bits,convert.is_signed,new_switch_expr.operand,**convert.tags,)else:new_switch_expr=ailment.Expr.Convert(None,convert.from_bits,convert.to_bits,convert.is_signed,new_switch_expr,**convert.tags)new_cases=OrderedDict()forcase_idx,case_nodeinnode.cases.items():ifisinstance(case_idx,int):new_cases[case_idx+v]=case_nodeelifisinstance(case_idx,tuple):new_cases[tuple(idx_+vforidx_incase_idx)]=case_nodeelse:raiseTypeError(f"Unsupported case_idx type {type(case_idx)}")new_node=SwitchCaseNode(new_switch_expr,new_cases,default_node=node.default_node,addr=node.addr)node=new_nodechanged=Truer=super()._handle_SwitchCase(node,**kwargs)ifchangedandrisNone:returnnodereturnr