fix: add store result inlining rule

bolt_expressions/node.py

@@ -56,6 +56,7 @@
     set_and_get_cleanup,
     set_to_self_removal,
     store_set_data_compare,
+    store_result_inlining,
 )
 from .typing import NbtTypeString
 from .casting import TypeCaster
@@ -264,6 +265,7 @@ def __init__(self, ctx: Context):
             set_and_get_cleanup=set_and_get_cleanup,
             noncommutative_set_collapsing=noncommutative_set_collapsing,
             commutative_set_collapsing=commutative_set_collapsing,
+            store_result_inlining=store_result_inlining,
             data_string_propagation=data_string_propagation,
             literal_to_constant_replacement=partial(
                 literal_to_constant_replacement, self.optimizer

bolt_expressions/optimizer.py

@@ -1550,3 +1550,44 @@ def convert_defined_boolean_condition(
             operands.append(operand)
 
         yield replace_operation(node, operands=tuple(operands))
+
+
+def store_result_inlining(nodes: Iterable[IrOperation]) -> Iterable[IrOperation]:
+    nodes = tuple(nodes)
+
+    defs = get_source_definitions(nodes)
+
+    stores: dict[int, list[IrStore]] = {}
+    removed: set[int] = set()
+
+    for i, node in enumerate(nodes):
+        if (
+            not isinstance(node, IrCast)
+            or not isinstance(node.right, IrSource)
+        ):
+            continue
+
+        source_def_i = get_reaching_definition(defs, node.right, i)
+
+        if source_def_i is None:
+            continue
+
+        if source_def_i != (i - 1):
+            continue
+
+        store = IrStore(
+            type=StoreType.result,
+            value=node.left,
+            cast_type=node.cast_type,
+            scale=node.scale
+        )
+        source_stores = stores.setdefault(source_def_i, [])
+        source_stores.append(store)
+        removed.add(i)
+
+    for i, node in enumerate(nodes):
+        if store := stores.get(i):
+            node = replace(node, store=IrChildren((*node.store, *store)))
+
+        if i not in removed:
+            yield node

tests/snapshots/examples__build_contrib_commands__0.data_pack/data/test/functions/main.mcfunction

@@ -20,5 +20,4 @@ scoreboard players add $foo obj.temp 10
 scoreboard players operation $i0 bolt.expr.temp = $delta obj.temp
 scoreboard players add $i0 bolt.expr.temp 10
 execute store result score $i1 bolt.expr.temp run data get storage example:main value 1
-scoreboard players operation $i0 bolt.expr.temp += $i1 bolt.expr.temp
-execute store result storage example:main value int 1 run scoreboard players get $i0 bolt.expr.temp
+execute store result storage example:main value int 1 run scoreboard players operation $i0 bolt.expr.temp += $i1 bolt.expr.temp

tests/snapshots/examples__build_operation_min_max__0.data_pack/data/test/functions/main.mcfunction

@@ -38,8 +38,7 @@ scoreboard players operation $f abc.main > $p abc.main
 execute store result score $i0 bolt.expr.temp run data get storage example:main a 1
 execute store result score $i1 bolt.expr.temp run data get storage example:main b 1
 scoreboard players operation $i0 bolt.expr.temp < $i1 bolt.expr.temp
-scoreboard players operation $i0 bolt.expr.temp < $a abc.main
-execute store result storage example:main value int 1 run scoreboard players get $i0 bolt.expr.temp
+execute store result storage example:main value int 1 run scoreboard players operation $i0 bolt.expr.temp < $a abc.main
 execute store result score $result abc.main run data get storage example:main b 1
 execute store result score $i0 bolt.expr.temp run data get storage example:main c 1
 scoreboard players operation $result abc.main < $i0 bolt.expr.temp

tests/snapshots/examples__build_operation_nbt_basics__0.data_pack/data/test/functions/literals.mcfunction

@@ -1,15 +1,12 @@
 data modify entity @s Health set value 10
 execute store result score $i0 bolt.expr.temp run data get entity @s Health 1
-scoreboard players add $i0 bolt.expr.temp 20
-execute store result entity @s Health int 1 run scoreboard players get $i0 bolt.expr.temp
+execute store result entity @s Health int 1 run scoreboard players add $i0 bolt.expr.temp 20
 execute store result score $i0 bolt.expr.temp run data get entity rx97 Health 1
-scoreboard players remove $i0 bolt.expr.temp 30
-execute store result entity rx97 Health int 1 run scoreboard players get $i0 bolt.expr.temp
+execute store result entity rx97 Health int 1 run scoreboard players remove $i0 bolt.expr.temp 30
 execute store result block 0 0 0 data.ur_mom int 40 run data get block 0 0 0 data.ur_mom 1
 execute store result block ~1 ~2 ~3 Items[0].Count double 0.02 run data get block ~1 ~2 ~3 Items[0].Count 1
 execute store result score $i0 bolt.expr.temp run data get entity TheWii Inventory[{Slot: -103b}].tag.data.bolt 1
-scoreboard players operation $i0 bolt.expr.temp %= $60 bolt.expr.const
-execute store result entity TheWii Inventory[{Slot: -103b}].tag.data.bolt int 1 run scoreboard players get $i0 bolt.expr.temp
+execute store result entity TheWii Inventory[{Slot: -103b}].tag.data.bolt int 1 run scoreboard players operation $i0 bolt.expr.temp %= $60 bolt.expr.const
 data modify entity @s Inventory[].tag.string set value "your mom"
 data modify entity @s Inventory[].tag.int set value 10
 data modify entity @s Inventory[].tag.float set value 10.0f

tests/snapshots/examples__build_operation_nbt_basics__0.data_pack/data/test/functions/main.mcfunction

@@ -1,6 +1,5 @@
 execute store result score @s abc.main run data get storage demo:prefix/temp items
-execute store result score $i0 bolt.expr.temp run data get storage demo:prefix/temp items
-execute store result storage demo:prefix/temp value short 1 run scoreboard players get $i0 bolt.expr.temp
+execute store result storage demo:prefix/temp value short 1 run data get storage demo:prefix/temp items
 execute store result score $x abc.main run data get storage demo:prefix/temp a
 execute store result score $i0 bolt.expr.temp run data get storage demo:prefix/temp b
 scoreboard players operation $x abc.main += $i0 bolt.expr.temp

tests/snapshots/examples__build_operation_nbt_basics__0.data_pack/data/test/functions/operations.mcfunction

@@ -17,25 +17,19 @@ scoreboard players operation @s abc.main = $i0 bolt.expr.temp
 data modify storage demo:prefix/temp foo set from storage demo:prefix/temp bar
 execute store result score $i0 bolt.expr.temp run data get storage demo:prefix/temp player 1
 execute store result score $i1 bolt.expr.temp run data get entity @s Health 1
-scoreboard players operation $i0 bolt.expr.temp += $i1 bolt.expr.temp
-execute store result storage demo:prefix/temp player int 1 run scoreboard players get $i0 bolt.expr.temp
+execute store result storage demo:prefix/temp player int 1 run scoreboard players operation $i0 bolt.expr.temp += $i1 bolt.expr.temp
 execute store result score $i0 bolt.expr.temp run data get storage demo:prefix/temp player 1
 execute store result score $i1 bolt.expr.temp run data get entity rx97 SelectedItem.Count 1
-scoreboard players operation $i0 bolt.expr.temp -= $i1 bolt.expr.temp
-execute store result storage demo:prefix/temp player int 1 run scoreboard players get $i0 bolt.expr.temp
+execute store result storage demo:prefix/temp player int 1 run scoreboard players operation $i0 bolt.expr.temp -= $i1 bolt.expr.temp
 execute store result score $i0 bolt.expr.temp run data get storage demo:prefix/temp player 1
 execute store result score $i1 bolt.expr.temp run data get entity TheWii SelectedItem.Damage 1
-scoreboard players operation $i0 bolt.expr.temp *= $i1 bolt.expr.temp
-execute store result storage demo:prefix/temp player int 1 run scoreboard players get $i0 bolt.expr.temp
+execute store result storage demo:prefix/temp player int 1 run scoreboard players operation $i0 bolt.expr.temp *= $i1 bolt.expr.temp
 execute store result score $i0 bolt.expr.temp run data get storage demo:prefix/temp player 1
 execute store result score $i1 bolt.expr.temp run data get block 0 0 0 RecordItem.data 1
-scoreboard players operation $i0 bolt.expr.temp /= $i1 bolt.expr.temp
-execute store result storage demo:prefix/temp player int 1 run scoreboard players get $i0 bolt.expr.temp
+execute store result storage demo:prefix/temp player int 1 run scoreboard players operation $i0 bolt.expr.temp /= $i1 bolt.expr.temp
 execute store result score $i0 bolt.expr.temp run data get storage demo:prefix/temp player 1
 execute store result score $i1 bolt.expr.temp run data get block ~1 ~2 ~3 Items[0] 1
-scoreboard players operation $i0 bolt.expr.temp %= $i1 bolt.expr.temp
-execute store result storage demo:prefix/temp player int 1 run scoreboard players get $i0 bolt.expr.temp
+execute store result storage demo:prefix/temp player int 1 run scoreboard players operation $i0 bolt.expr.temp %= $i1 bolt.expr.temp
 execute store result score $i0 bolt.expr.temp run data get block ~1 ~2 ~3 Items[-1].Count 1
 execute store result score $i1 bolt.expr.temp run data get entity @s SelectedItem.Count 1
-scoreboard players operation $i0 bolt.expr.temp += $i1 bolt.expr.temp
-execute store result block ~1 ~2 ~3 Items[-1].Count int 1 run scoreboard players get $i0 bolt.expr.temp
+execute store result block ~1 ~2 ~3 Items[-1].Count int 1 run scoreboard players operation $i0 bolt.expr.temp += $i1 bolt.expr.temp

tests/snapshots/examples__build_operation_nbt_scale__0.data_pack/data/test/functions/main.mcfunction

@@ -8,20 +8,16 @@ say # get and set scaled data
 execute store result storage demo out double 1.0 run data get storage demo value 1
 say # set scaled score
 scoreboard players operation $i0 bolt.expr.temp = $value obj
-scoreboard players add $i0 bolt.expr.temp 1
-execute store result storage demo out int 50 run scoreboard players get $i0 bolt.expr.temp
+execute store result storage demo out int 50 run scoreboard players add $i0 bolt.expr.temp 1
 say # set scaled score
 scoreboard players operation $i0 bolt.expr.temp = $value obj
-scoreboard players add $i0 bolt.expr.temp 1
-execute store result storage demo out double 0.1 run scoreboard players get $i0 bolt.expr.temp
+execute store result storage demo out double 0.1 run scoreboard players add $i0 bolt.expr.temp 1
 say # get scaled data
 execute store result score $i0 bolt.expr.temp run data get storage demo value 100
-scoreboard players add $i0 bolt.expr.temp 3
-execute store result storage demo out int 1 run scoreboard players get $i0 bolt.expr.temp
+execute store result storage demo out int 1 run scoreboard players add $i0 bolt.expr.temp 3
 say # get scaled data, operate and set it back scaled down
 execute store result score $i0 bolt.expr.temp run data get storage demo value 100
-scoreboard players add $i0 bolt.expr.temp 5
-execute store result storage demo out double 0.01 run scoreboard players get $i0 bolt.expr.temp
+execute store result storage demo out double 0.01 run scoreboard players add $i0 bolt.expr.temp 5
 say # get scaled data
 execute store result score #offset obj run data get entity @s Motion[0] 100
 scoreboard players operation #offset obj -= #x obj

tests/snapshots/examples__build_operation_nbt_type__0.data_pack/data/test/functions/main.mcfunction

@@ -7,22 +7,18 @@ data modify storage demo n set from storage bolt.expr:temp 2384k242hd495_4
 execute store result storage bolt.expr:temp 2384k242hd495_5 float 1 run data get storage bolt.expr:temp 2384k242hd495_4 1
 data modify storage demo m set from storage bolt.expr:temp 2384k242hd495_5
 execute store result score $i0 bolt.expr.temp run data get storage demo x 100
-scoreboard players add $i0 bolt.expr.temp 1
-execute store result storage bolt.expr:temp 2384k242hd495_12 double 0.01 run scoreboard players get $i0 bolt.expr.temp
+execute store result storage bolt.expr:temp 2384k242hd495_12 double 0.01 run scoreboard players add $i0 bolt.expr.temp 1
 data modify storage demo x set from storage bolt.expr:temp 2384k242hd495_12
 execute store result score $i0 bolt.expr.temp run data get storage demo num 100
-scoreboard players add $i0 bolt.expr.temp 1
-execute store result storage bolt.expr:temp 2384k242hd495_19 double 1 run scoreboard players get $i0 bolt.expr.temp
+execute store result storage bolt.expr:temp 2384k242hd495_19 double 1 run scoreboard players add $i0 bolt.expr.temp 1
 execute store result storage demo num double 0.01 run data get storage bolt.expr:temp 2384k242hd495_19 1
 execute store result storage bolt.expr:temp 2384k242hd495_24 short 1 run scoreboard players get $val obj
 execute store result storage demo a int 100 run data get storage bolt.expr:temp 2384k242hd495_24 1
 execute store result storage demo a double 0.1 run data get storage bolt.expr:temp 2384k242hd495_24 1
 scoreboard players operation $i0 bolt.expr.temp = $a obj
-scoreboard players add $i0 bolt.expr.temp 1
-execute store result storage demo a int 1 run scoreboard players get $i0 bolt.expr.temp
+execute store result storage demo a int 1 run scoreboard players add $i0 bolt.expr.temp 1
 scoreboard players operation $i0 bolt.expr.temp = $b obj
-scoreboard players add $i0 bolt.expr.temp 2
-execute store result storage bolt.expr:temp 2384k242hd495_32 double 1 run scoreboard players get $i0 bolt.expr.temp
+execute store result storage bolt.expr:temp 2384k242hd495_32 double 1 run scoreboard players add $i0 bolt.expr.temp 2
 execute store result storage demo a double 0.1 run data get storage bolt.expr:temp 2384k242hd495_32 1
 say nicer
 execute store result storage bolt.expr:temp 2384k242hd495_37 double 0.01 run scoreboard players get $foo obj
@@ -42,8 +38,7 @@ execute store result storage demo value double 0.1 run scoreboard players get @s
 execute store result storage demo value byte 1 run scoreboard players get @s obj
 execute store result score $i0 bolt.expr.temp run data get storage demo x1 100
 execute store result score $i1 bolt.expr.temp run data get storage demo x2 100
-scoreboard players operation $i0 bolt.expr.temp += $i1 bolt.expr.temp
-execute store result storage demo value double 0.005 run scoreboard players get $i0 bolt.expr.temp
+execute store result storage demo value double 0.005 run scoreboard players operation $i0 bolt.expr.temp += $i1 bolt.expr.temp
 say casting to byte before appending to list
 execute store result storage demo value byte 1 run scoreboard players get $value obj
 data modify storage demo list append from storage demo value

tests/snapshots/examples__build_source_typing__0.data_pack/data/test/functions/main.mcfunction

@@ -57,8 +57,7 @@ data modify storage name:path value prepend value 0.0d
 data modify storage name:path value insert 4 value 13.0d
 data modify storage bolt.expr:temp i0 set value {id: "bla", Count: 5, tag: {CustomModelData: 14.2f, Items: [{id: "stone", Count: 3}]}}
 data modify storage bolt.expr:temp i0 merge from storage name:path other_item
-data modify storage bolt.expr:temp i0 merge value {tag: {Items: [{Count: 23}]}}
-execute store result storage name:path new_item int 1 run data get storage bolt.expr:temp i0 1
+execute store result storage name:path new_item int 1 run data modify storage bolt.expr:temp i0 merge value {tag: {Items: [{Count: 23}]}}
 data modify storage name:path name set value "george"
 data modify storage name:path flag set value 127b
 data modify storage name:path flags set value [B; 1B, 2B, 3B, 127B]