[Lang] Support chained assignments

[gaocegege]

Fix #2659

This PR is to support chained assignments

Test case:

@ti.kernel
def func():
    (a, b) = (c, d) = (7, 8)
    aa = bb = -1
    print(a, b, c, d)
    print(aa, bb)

func()

Python generated code

Initial AST:
def func():
    a, b = c, d = 7, 8
    aa = bb = -1
    print(a, b, c, d)
    print(aa, bb)

Preprocessed:
def func():
    import taichi as ti
    if 1:
        if 1:
            __tmp_tuple = ti.expr_init_list((7, 8), 2)
            a = ti.expr_init(__tmp_tuple[0])
            b = ti.expr_init(__tmp_tuple[1])
            del __tmp_tuple
        if 1:
            __tmp_tuple = ti.expr_init_list((7, 8), 2)
            c = ti.expr_init(__tmp_tuple[0])
            d = ti.expr_init(__tmp_tuple[1])
            del __tmp_tuple
    if 1:
        aa = bb = ti.expr_init(-1)
        aa = bb = ti.expr_init(-1)
    ti.ti_print(a, b, c, d)
    ti.ti_print(aa, bb)
    del bb
    del aa
    del d
    del c
    del b
    del a

Result:

7 8 7 8
-1 -1

[netlify]

✔️ Deploy Preview for jovial-fermat-aa59dc ready!

🔨 Explore the source changes: 8b42b67

🔍 Inspect the deploy log: https://app.netlify.com/sites/jovial-fermat-aa59dc/deploys/615fe7734e0a8b00089a97bd

😎 Browse the preview: https://deploy-preview-3062--jovial-fermat-aa59dc.netlify.app

[gaocegege]

/cc @k-ye

Do we have the unit tests or integration tests for the Python code transformation logic?

[k-ye]

Yep, they're inside https://github.com/taichi-dev/taichi/tree/master/tests/python. See also https://docs.taichi.graphics/docs/lang/articles/contribution/write_test :-)

[gaocegege]

Yep, they're inside https://github.com/taichi-dev/taichi/tree/master/tests/python. See also https://docs.taichi.graphics/docs/lang/articles/contribution/write_test :-)

Thanks. I will add more test cases for the feature.

[gaocegege]

I added some test cases (unpack, chained, chained + unpack, assign) for the feature, and the PR is ready to review, I think.

Please take a look, thanks for your time 😄

[gaocegege]

Please take a look, thanks for your time 😄

[lin-hitonami]

taichi/python/taichi/lang/stmt_builder.py

Lines 118 to 145 in fe5b54d

	is_local = isinstance(node_target, ast.Name)
	if is_local and ctx.is_creation(node_target.id):
	var_name = node_target.id
	# Create, no AST resolution needed
	init = ast.Attribute(value=ast.Name(id='ti',
	ctx=ast.Load()),
	attr='expr_init',
	ctx=ast.Load())
	rhs = ast.Call(
	func=init,
	args=[node.value],
	keywords=[],
	)
	ctx.create_variable(var_name)
	assign_stmts.append(
	ast.copy_location(
	ast.Assign(targets=node.targets,
	value=rhs,
	type_comment=None), node))
	else:
	# Assign
	node_target.ctx = ast.Load()
	func = ast.Attribute(value=node_target,
	attr='assign',
	ctx=ast.Load())
	call = ast.Call(func=func, args=[node.value], keywords=[])
	assign_stmts.append(
	ast.copy_location(ast.Expr(value=call), node))

Can we make this part a static method and use it in both tuple and non-tuple cases?

[gaocegege]

taichi/python/taichi/lang/stmt_builder.py

Lines 118 to 145 in fe5b54d

	is_local = isinstance(node_target, ast.Name)
	if is_local and ctx.is_creation(node_target.id):
	var_name = node_target.id
	# Create, no AST resolution needed
	init = ast.Attribute(value=ast.Name(id='ti',
	ctx=ast.Load()),
	attr='expr_init',
	ctx=ast.Load())
	rhs = ast.Call(
	func=init,
	args=[node.value],
	keywords=[],
	)
	ctx.create_variable(var_name)
	assign_stmts.append(
	ast.copy_location(
	ast.Assign(targets=node.targets,
	value=rhs,
	type_comment=None), node))
	else:
	# Assign
	node_target.ctx = ast.Load()
	func = ast.Attribute(value=node_target,
	attr='assign',
	ctx=ast.Load())
	call = ast.Call(func=func, args=[node.value], keywords=[])
	assign_stmts.append(
	ast.copy_location(ast.Expr(value=call), node))

Can we make this part a static method and use it in both tuple and non-tuple cases?

Do you mean that we do not make tuple case a func call, we just wrap the detailed logic (assign and var creation) into one func, then use it in tuple and non-tuple?

[lin-hitonami]

taichi/python/taichi/lang/stmt_builder.py

Lines 118 to 145 in fe5b54d

	is_local = isinstance(node_target, ast.Name)
	if is_local and ctx.is_creation(node_target.id):
	var_name = node_target.id
	# Create, no AST resolution needed
	init = ast.Attribute(value=ast.Name(id='ti',
	ctx=ast.Load()),
	attr='expr_init',
	ctx=ast.Load())
	rhs = ast.Call(
	func=init,
	args=[node.value],
	keywords=[],
	)
	ctx.create_variable(var_name)
	assign_stmts.append(
	ast.copy_location(
	ast.Assign(targets=node.targets,
	value=rhs,
	type_comment=None), node))
	else:
	# Assign
	node_target.ctx = ast.Load()
	func = ast.Attribute(value=node_target,
	attr='assign',
	ctx=ast.Load())
	call = ast.Call(func=func, args=[node.value], keywords=[])
	assign_stmts.append(
	ast.copy_location(ast.Expr(value=call), node))

Can we make this part a static method and use it in both tuple and non-tuple cases?

Do you mean that we do not make tuple case a func call, we just wrap the detailed logic (assign and var creation) into one func, then use it in tuple and non-tuple?

I just think the logic of assign and var creation is a bit duplicated. If you think the function is too long, you can also split it into small functions. (just my opinion)

[gaocegege]

Do you mean that we do not make tuple case a func call, we just wrap the detailed logic (assign and var creation) into one func, then use it in tuple and non-tuple?

SGTM. It makes sense.

[lin-hitonami]

LGTM!