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(']')
]))
]