diff --git a/src/imp/gen_map_get.rs b/src/imp/gen_map_get.rs index 0139b097..736bb63d 100644 --- a/src/imp/gen_map_get.rs +++ b/src/imp/gen_map_get.rs @@ -15,14 +15,22 @@ impl Definition { impl Stmt { fn gen_map_get(&mut self, id: &mut usize) { match self { - Stmt::Assign { pat: _, val, nxt } => { + Stmt::Assign { pat, val, nxt } => { + let key_substitutions = + if let AssignPattern::MapSet(_, key) = pat { key.substitute_map_gets(id) } else { Vec::new() }; + if let Some(nxt) = nxt { nxt.gen_map_get(id); } + let substitutions = val.substitute_map_gets(id); if !substitutions.is_empty() { *self = gen_get(self, substitutions); } + + if !key_substitutions.is_empty() { + *self = gen_get(self, key_substitutions); + } } Stmt::Ask { pat: _, val, nxt } => { nxt.gen_map_get(id); diff --git a/src/imp/mod.rs b/src/imp/mod.rs index f77120ef..70d7b82c 100644 --- a/src/imp/mod.rs +++ b/src/imp/mod.rs @@ -46,9 +46,10 @@ pub struct MatchArm { pub rgt: Stmt, } -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Default)] pub enum AssignPattern { // "*" + #[default] Eraser, // [a-zA-Z_]+ Var(Name), diff --git a/tests/golden_tests/run_file/nested_map_set.bend b/tests/golden_tests/run_file/nested_map_set.bend new file mode 100644 index 00000000..6d8f56de --- /dev/null +++ b/tests/golden_tests/run_file/nested_map_set.bend @@ -0,0 +1,5 @@ +def main(): + map = { 0: 1, 1: 10, 2: 0, 3: 1, 4: 3 } + map[map[0]] = 99 + map[map[2]] = 1 + return map[map[map[4]]] + map[map[2]] diff --git a/tests/snapshots/run_file__nested_map_set.bend.snap b/tests/snapshots/run_file__nested_map_set.bend.snap new file mode 100644 index 00000000..ea6941a0 --- /dev/null +++ b/tests/snapshots/run_file__nested_map_set.bend.snap @@ -0,0 +1,9 @@ +--- +source: tests/golden_tests.rs +input_file: tests/golden_tests/run_file/nested_map_set.bend +--- +NumScott: +100 + +Scott: +100