RakuAST: Create optimized code for array slot assignment

When assigning into an array slot, we can directly call postcircumfix:<[ ]> with an additional argument.
rakudo · Oct 28, 2022 · e5cfb41 · e5cfb41
1 parent 31f28de
commit e5cfb41
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 1 deletion.
diff --git a/src/Raku/Actions.nqp b/src/Raku/Actions.nqp
@@ -485,12 +485,23 @@ class Raku::Actions is HLL::Actions does Raku::CommonActions {
     ## Expression parsing and operators
     ##
 
+    my %infix_specials := nqp::hash(
+        '=', -> $actions, $/, $sym { assign_op($actions, $/, $/[0].ast, $/[1].ast) },
+    );
+
     method EXPR($/, $KEY?) {
         my $ast := $/.ast // $<OPER>.ast;
         if $KEY {
             my $key := nqp::lc($KEY);
             if $KEY eq 'INFIX' {
                 my $sym := $<infix><sym>;
+                if nqp::existskey(%infix_specials, $sym) {
+                    my $ast := %infix_specials{$sym}(self, $/, $sym);
+                    if $ast {
+                        self.attach: $/, $ast;
+                        return;
+                    }
+                }
                 if $sym && $sym eq '??' {
                     self.attach: $/, self.r('Ternary').new:
                         condition => $/[0].ast,
@@ -546,6 +557,18 @@ class Raku::Actions is HLL::Actions does Raku::CommonActions {
         }
     }
 
+    sub assign_op($actions, $/, $lhs_ast, $rhs_ast, :$initialize) {
+        if (
+            nqp::istype($lhs_ast, $actions.r('ApplyPostfix'))
+            && nqp::istype($lhs_ast.postfix, $actions.r('Postcircumfix', 'ArrayIndex'))
+            && nqp::elems($lhs_ast.IMPL-UNWRAP-LIST($lhs_ast.postfix.index.statements)) == 1
+        ) {
+            $lhs_ast.postfix.set-assignee($rhs_ast);
+            return $lhs_ast;
+        }
+        return;
+    }
+
     method prefixish($/) {
         my $ast := $<OPER>.ast // self.r('Prefix').new(~$<prefix><sym>);
         if $<prefix_postfix_meta_operator> {

diff --git a/src/Raku/ast/expressions.rakumod b/src/Raku/ast/expressions.rakumod
@@ -922,14 +922,20 @@ class RakuAST::Postcircumfix is RakuAST::Postfixish {
 # A postcircumfix array index operator, possibly multi-dimensional.
 class RakuAST::Postcircumfix::ArrayIndex is RakuAST::Postcircumfix is RakuAST::Lookup {
     has RakuAST::SemiList $.index;
+    has RakuAST::Expression $.assignee;
 
-    method new(RakuAST::SemiList $index) {
+    method new(RakuAST::SemiList $index, RakuAST::Expression :$assignee) {
         my $obj := nqp::create(self);
         nqp::bindattr($obj, RakuAST::Postcircumfix::ArrayIndex, '$!index', $index);
+        nqp::bindattr($obj, RakuAST::Postcircumfix::ArrayIndex, '$!assignee', $assignee // RakuAST::Expression);
         nqp::bindattr($obj, RakuAST::Postfixish, '$!colonpairs', []);
         $obj
     }
 
+    method set-assignee(RakuAST::Expression $assignee) {
+        nqp::bindattr(self, RakuAST::Postcircumfix::ArrayIndex, '$!assignee', $assignee);
+    }
+
     method can-be-bound-to() {
         True
     }
@@ -947,13 +953,15 @@ class RakuAST::Postcircumfix::ArrayIndex is RakuAST::Postcircumfix is RakuAST::L
 
     method visit-children(Code $visitor) {
         $visitor($!index);
+        $visitor($!assignee) if $!assignee;
         self.visit-colonpairs($visitor);
     }
 
     method IMPL-POSTFIX-QAST(RakuAST::IMPL::QASTContext $context, Mu $operand-qast) {
         my $name := self.resolution.lexical-name;
         my $op := QAST::Op.new( :op('call'), :$name, $operand-qast );
         $op.push($!index.IMPL-TO-QAST($context)) unless $!index.is-empty;
+        $op.push($!assignee.IMPL-TO-QAST($context)) if $!assignee;
         $op
     }