def _swap_pp_rule(eqn, context, settings): y, = eqn.outvars x, v, *idx = eqn.invars idx = ','.join(core.pp_var(i, context) for i in idx) if type(y) is core.DropVar: # In the case of a set (ignored return value), # pretty print `_ = swap x v i` as `x[i] <- v` del y return [ pp_ref( pp.concat([ pp.text(core.pp_var(x, context)), pp.text('['), pp.text(idx), pp.text(']') ])), pp.text(' <- '), pp.text(core.pp_var(v, context)) ] else: # pretty-print `y:T = swap x v i` as `y:T, x[i] <- x[i], v` x_i = pp.concat([ pp.text(core.pp_var(x, context)), pp.text('['), pp.text(idx), pp.text(']') ]) y = core.pp_vars([y], context, print_shapes=settings.print_shapes) return [ y, pp.text(', '), x_i, pp.text(' <- '), x_i, pp.text(', '), pp.text(core.pp_var(v, context)) ]
def _addupdate_pp_rule(eqn, context, settings): # pretty-print ` = addupdate x i v` as `x[i] += v` () = eqn.outvars x, v, *idx = eqn.invars idx = ','.join(core.pp_var(i, context) for i in idx) return [ pp_ref( pp.concat([ pp.text(core.pp_var(x, context)), pp.text('['), pp.text(idx), pp.text(']') ])), pp.text(' += '), pp.text(core.pp_var(v, context)) ]
def _get_pp_rule(eqn, context, settings): # Pretty prints `a = get x i` as `x[i] <- a` y, = eqn.outvars x, *idx = eqn.invars idx = ','.join(core.pp_var(i, context) for i in idx) lhs = core.pp_vars([y], context, print_shapes=settings.print_shapes) return [ lhs, pp.text(' <- '), pp_ref( pp.concat([ pp.text(core.pp_var(x, context)), pp.text('['), pp.text(idx), pp.text(']') ])) ]