[Misc] Refactored flattend_values() to avoid potential conflicts in f…

…lattened statements (#6749)

Issue: #5819

Overriding the flattened statement `stmt` of an `Expression` can cause
conflicts, for example:
```
@ti.kernel
def test():
    x = ti.Vector([1, 2, 3, 4])
    tmp = x + x[0] # implicit broadcast
```

In `x + x[0]`, the `x` on the lhs serves as rvalue whereas the `x` in
the `x[0]` serves as a lvalue, so the result of `flatten_rvalue()` and
`flatten_lvalue()` will override each other.

To avoid such conflicts, this PR refactored the `flatten_values()`
functions:
1. Flattened statement `stmt` of an `Expression` will only get modified
by `Expression::flatten()`, any other overriding will be forbidden.
2. `flatten_rvalue()` and `flatten_lvalue()` now returns the flattened
statement as the result. External users such as `irpass::lower_ast()`
will turn to use the returned statement.

Co-authored-by: Yi Xu <xy_xuyi@foxmail.com>

taichi/analysis/gen_offline_cache_key.cpp

@@ -571,7 +571,7 @@ class ASTSerializer : public IRVisitor, public ExpressionVisitor {
       emit(expr.const_value);
       emit(expr.atomic);
       auto *e = expr.expr.get();
-      emit(e->stmt);
+      emit(e->get_flattened_stmt());
       emit(e->attributes);
       emit(e->ret_type);
       expr.expr->accept(this);

taichi/ir/expression.h

@@ -11,8 +11,10 @@ class ExpressionVisitor;
 
 // always a tree - used as rvalues
 class Expression {
- public:
+ protected:
   Stmt *stmt;
+
+ public:
   std::string tb;
   std::map<std::string, std::string> attributes;
   DataType ret_type;
@@ -53,6 +55,10 @@ class Expression {
 
   virtual ~Expression() {
   }
+
+  Stmt *get_flattened_stmt() const {
+    return stmt;
+  }
 };
 
 class ExprGroup {