Overloaded indexed assignments

Author: Jorge Aparicio

src/libcollections/btree/map.rs

@@ -22,6 +22,8 @@ use core::fmt::Debug;
 use core::hash::{Hash, Hasher};
 use core::iter::{Map, FromIterator};
 use core::ops::Index;
+#[cfg(not(stage0))]
+use core::ops::{IndexAssign, IndexMut};
 use core::{fmt, mem, usize};
 use Bound::{self, Included, Excluded, Unbounded};
 
@@ -949,6 +951,28 @@ impl<'a, K: Ord, Q: ?Sized, V> Index<&'a Q> for BTreeMap<K, V>
     }
 }
 
+#[cfg(not(stage0))]
+// TODO(japaric) update issue number
+#[unstable(feature = "index_assign_trait", reason="recently added", issue="0")]
+impl<'a, K: Ord, Q: ?Sized, V> IndexMut<&'a Q> for BTreeMap<K, V>
+    where K: Borrow<Q>, Q: Ord
+{
+    #[inline]
+    fn index_mut(&mut self, key: &Q) -> &mut V {
+        self.get_mut(key).expect("no entry found for key")
+    }
+}
+
+#[cfg(not(stage0))]
+// TODO(japaric) update issue number
+#[unstable(feature = "index_assign_trait", reason="recently added", issue="0")]
+impl<'a, K: Ord, V> IndexAssign<K, V> for BTreeMap<K, V> {
+    #[inline]
+    fn index_assign(&mut self, key: K, value: V) {
+        self.insert(key, value);
+    }
+}
+
 /// Genericises over how to get the correct type of iterator from the correct type
 /// of Node ownership.
 trait Traverse<N> {

src/libcollections/lib.rs

@@ -59,6 +59,7 @@
 #![feature(decode_utf16)]
 #![feature(utf8_error)]
 #![cfg_attr(test, feature(rand, test))]
+#![cfg_attr(not(stage0), feature(index_assign_trait))]
 
 #![feature(no_std)]
 #![no_std]

src/libcore/ops.rs

@@ -1539,6 +1539,44 @@ pub trait IndexMut<Idx: ?Sized>: Index<Idx> {
     fn index_mut(&mut self, index: Idx) -> &mut Self::Output;
 }
 
+/// The `IndexAssign` trait is used to specify the functionality of indexed assignment
+/// operations like `arr[idx] = rhs`.
+///
+/// # Examples
+///
+/// A trivial implementation of `IndexAssign`. When `Foo[Bar] = Baz` happens, it ends up
+/// calling `index_assign`, and therefore, `main` prints `Indexed assignment!`.
+///
+/// ```
+/// #![feature(indexed_assignments)]
+/// #![feature(index_assign_trait)]
+///
+/// use std::ops::{IndexAssign};
+///
+/// #[derive(Copy, Clone)]
+/// struct Foo;
+/// struct Bar;
+/// struct Baz;
+///
+/// impl IndexAssign<Bar, Baz> for Foo {
+///     fn index_assign(&mut self, _index: Bar, _rhs: Baz) {
+///         println!("Indexed assignment!");
+///     }
+/// }
+///
+/// fn main() {
+///     Foo[Bar] = Baz;
+/// }
+/// ```
+#[cfg(not(stage0))]
+#[lang = "index_assign"]
+// TODO(japaric) update issue number
+#[unstable(feature = "index_assign_trait", reason = "recently added", issue = "0")]
+pub trait IndexAssign<Idx, Rhs> {
+    /// The method for the indexed assignment (`Foo[Bar] = Baz`) operation
+    fn index_assign(&mut self, index: Idx, rhs: Rhs);
+}
+
 /// An unbounded range.
 #[derive(Copy, Clone, PartialEq, Eq)]
 #[lang = "range_full"]

src/librustc/middle/expr_use_visitor.rs

@@ -517,8 +517,16 @@ impl<'d,'t,'a,'tcx> ExprUseVisitor<'d,'t,'a,'tcx> {
             }
 
             hir::ExprAssign(ref lhs, ref rhs) => {
-                self.mutate_expr(expr, &**lhs, JustWrite);
-                self.consume_expr(&**rhs);
+                if let hir::ExprIndex(ref base, ref idx) = lhs.node {
+                    if !self.walk_overloaded_operator(expr, base, vec![idx, rhs],
+                                                      PassArgs::ByValue) {
+                        self.mutate_expr(expr, &**lhs, JustWrite);
+                        self.consume_expr(&**rhs);
+                    }
+                } else {
+                    self.mutate_expr(expr, &**lhs, JustWrite);
+                    self.consume_expr(&**rhs);
+                }
             }
 
             hir::ExprCast(ref base, _) => {

src/librustc/middle/lang_items.rs

@@ -298,6 +298,7 @@ lets_do_this! {
     ShrAssignTraitLangItem,          "shr_assign",              shr_assign_trait;
     IndexTraitLangItem,              "index",                   index_trait;
     IndexMutTraitLangItem,           "index_mut",               index_mut_trait;
+    IndexAssignTraitLangItem,        "index_assign",            index_assign_trait;
     RangeStructLangItem,             "range",                   range_struct;
     RangeFromStructLangItem,         "range_from",              range_from_struct;
     RangeToStructLangItem,           "range_to",                range_to_struct;

src/librustc_trans/trans/callee.rs

@@ -868,6 +868,8 @@ pub enum CallArgs<'a, 'tcx> {
     // arguments should be auto-referenced
     ArgOverloadedOp(Datum<'tcx, Expr>, Option<(Datum<'tcx, Expr>, ast::NodeId)>, bool),
 
+    ArgIndexedAssignment(Vec<Datum<'tcx, Expr>>),
+
     // Supply value of arguments as a list of expressions that must be
     // translated, for overloaded call operators.
     ArgOverloadedCall(Vec<&'a hir::Expr>),
@@ -1055,6 +1057,15 @@ pub fn trans_args<'a, 'blk, 'tcx>(cx: Block<'blk, 'tcx>,
                 assert_eq!(arg_tys.len(), 1);
             }
         }
+        ArgIndexedAssignment(args) => {
+            assert!(!variadic);
+            assert_eq!(arg_tys.len(), 3);
+
+            for (i, arg) in args.into_iter().enumerate() {
+                bcx = trans_arg_datum(bcx, arg_tys[i], arg, arg_cleanup_scope, DontAutorefArg,
+                                      llargs);
+            }
+        }
         ArgVals(vs) => {
             llargs.push_all(vs);
         }

src/librustc_trans/trans/expr.rs

@@ -999,6 +999,32 @@ fn trans_rvalue_stmt_unadjusted<'blk, 'tcx>(bcx: Block<'blk, 'tcx>,
         hir::ExprLoop(ref body, _) => {
             controlflow::trans_loop(bcx, expr, &**body)
         }
+        // indexed assignments `a[b] = c`
+        hir::ExprAssign(ref lhs, ref rhs) if bcx.tcx().is_method_call(expr.id) => {
+            if let hir::ExprIndex(ref base, ref idx) = lhs.node {
+                let rhs = unpack_datum!(bcx, trans(bcx, &**rhs));
+
+                // NOTE `lhs` is not an lvalue in this case, but we still need to create a cleanup
+                // scope that covers `base` and `idx`
+                let cleanup_debug_loc = debuginfo::get_cleanup_debug_loc_for_ast_node(bcx.ccx(),
+                                                                                      lhs.id,
+                                                                                      lhs.span,
+                                                                                      false);
+
+                bcx.fcx.push_ast_cleanup_scope(cleanup_debug_loc);
+
+                let base = unpack_datum!(bcx, trans(bcx, &**base));
+                let idx = unpack_datum!(bcx, trans(bcx, &**idx));
+
+                bcx.fcx.pop_and_trans_ast_cleanup_scope(bcx, lhs.id);
+
+                trans_indexed_assignment(bcx, expr, MethodCall::expr(expr.id), base, idx, rhs).bcx
+            } else {
+                bcx.tcx().sess.span_bug(
+                    expr.span,
+                    "the only overloadable assignments are indexed assignments");
+            }
+        },
         hir::ExprAssign(ref dst, ref src) => {
             let src_datum = unpack_datum!(bcx, trans(bcx, &**src));
             let dst_datum = unpack_datum!(bcx, trans_to_lvalue(bcx, &**dst, "assign"));
@@ -2031,6 +2057,25 @@ fn trans_overloaded_op<'blk, 'tcx>(bcx: Block<'blk, 'tcx>,
                              dest)
 }
 
+fn trans_indexed_assignment<'blk, 'tcx>(bcx: Block<'blk, 'tcx>,
+                                        expr: &hir::Expr,
+                                        method_call: MethodCall,
+                                        base: Datum<'tcx, Expr>,
+                                        idx: Datum<'tcx, Expr>,
+                                        rhs: Datum<'tcx, Expr>)
+                                   -> Result<'blk, 'tcx> {
+    callee::trans_call_inner(bcx,
+                             expr.debug_loc(),
+                             |bcx, arg_cleanup_scope| {
+                                meth::trans_method_callee(bcx,
+                                                          method_call,
+                                                          None,
+                                                          arg_cleanup_scope)
+                             },
+                             callee::ArgIndexedAssignment(vec![base, idx, rhs]),
+                             None)
+}
+
 fn trans_overloaded_call<'a, 'blk, 'tcx>(mut bcx: Block<'blk, 'tcx>,
                                          expr: &hir::Expr,
                                          callee: &'a hir::Expr,
@@ -2676,8 +2721,8 @@ fn expr_kind(tcx: &ty::ctxt, expr: &hir::Expr) -> ExprKind {
         // Overloaded operations are generally calls, and hence they are
         // generated via DPS, but there are a few exceptions:
         return match expr.node {
-            // `a += b` has a unit result.
-            hir::ExprAssignOp(..) => ExprKind::RvalueStmt,
+            // `a += b` and `a[b] = c` have a unit result.
+            hir::ExprAssignOp(..) | hir::ExprAssign(..) => ExprKind::RvalueStmt,
 
             // the deref method invoked for `*a` always yields an `&T`
             hir::ExprUnary(hir::UnDeref, _) => ExprKind::Lvalue,

src/librustc_typeck/check/method/confirm.rs

@@ -574,7 +574,8 @@ impl<'a,'tcx> ConfirmContext<'a,'tcx> {
                             autoderefs,
                             unsize,
                             PreferMutLvalue,
-                            index_expr_ty);
+                            index_expr_ty,
+                            false);
 
                         if let Some((input_ty, return_ty)) = result {
                             demand::suptype(self.fcx, index_expr.span, input_ty, index_expr_ty);

src/librustc_typeck/check/mod.rs

@@ -2247,7 +2247,8 @@ fn lookup_indexing<'a, 'tcx>(fcx: &FnCtxt<'a, 'tcx>,
                              base_expr: &'tcx hir::Expr,
                              base_ty: Ty<'tcx>,
                              idx_ty: Ty<'tcx>,
-                             lvalue_pref: LvaluePreference)
+                             lvalue_pref: LvaluePreference,
+                             is_assignment: bool)
                              -> Option<(/*index type*/ Ty<'tcx>, /*element type*/ Ty<'tcx>)>
 {
     // FIXME(#18741) -- this is almost but not quite the same as the
@@ -2262,7 +2263,7 @@ fn lookup_indexing<'a, 'tcx>(fcx: &FnCtxt<'a, 'tcx>,
                                                lvalue_pref,
                                                |adj_ty, idx| {
         try_index_step(fcx, MethodCall::expr(expr.id), expr, base_expr,
-                       adj_ty, idx, false, lvalue_pref, idx_ty)
+                       adj_ty, idx, false, lvalue_pref, idx_ty, is_assignment)
     });
 
     if final_mt.is_some() {
@@ -2274,7 +2275,7 @@ fn lookup_indexing<'a, 'tcx>(fcx: &FnCtxt<'a, 'tcx>,
     if let ty::TyArray(element_ty, _) = ty.sty {
         let adjusted_ty = fcx.tcx().mk_slice(element_ty);
         try_index_step(fcx, MethodCall::expr(expr.id), expr, base_expr,
-                       adjusted_ty, autoderefs, true, lvalue_pref, idx_ty)
+                       adjusted_ty, autoderefs, true, lvalue_pref, idx_ty, is_assignment)
     } else {
         None
     }
@@ -2292,21 +2293,45 @@ fn try_index_step<'a, 'tcx>(fcx: &FnCtxt<'a, 'tcx>,
                             autoderefs: usize,
                             unsize: bool,
                             lvalue_pref: LvaluePreference,
-                            index_ty: Ty<'tcx>)
+                            index_ty: Ty<'tcx>,
+                            is_assignment: bool)
                             -> Option<(/*index type*/ Ty<'tcx>, /*element type*/ Ty<'tcx>)>
 {
     let tcx = fcx.tcx();
     debug!("try_index_step(expr={:?}, base_expr.id={:?}, adjusted_ty={:?}, \
-                           autoderefs={}, unsize={}, index_ty={:?})",
+                           autoderefs={}, unsize={}, index_ty={:?}, is_assignment={})",
            expr,
            base_expr,
            adjusted_ty,
            autoderefs,
            unsize,
-           index_ty);
+           index_ty,
+           is_assignment);
 
     let input_ty = fcx.infcx().next_ty_var();
 
+    // Try `IndexAssign` if this is an assignment operation
+    if is_assignment {
+        return tcx.lang_items.index_assign_trait().and_then(|trait_did| {
+            let rhs_ty = fcx.infcx().next_ty_var();
+
+            method::lookup_in_trait_adjusted(fcx,
+                                             expr.span,
+                                             Some(&*base_expr),
+                                             token::intern("index_assign"),
+                                             trait_did,
+                                             autoderefs,
+                                             unsize,
+                                             adjusted_ty,
+                                             Some(vec![input_ty, rhs_ty]))
+                .map(|method| {
+                    fcx.inh.tables.borrow_mut().method_map.insert(method_call, method);
+
+                    (input_ty, rhs_ty)
+                })
+        })
+    }
+
     // First, try built-in indexing.
     match (adjusted_ty.builtin_index(), &index_ty.sty) {
         (Some(ty), &ty::TyUint(ast::TyUs)) | (Some(ty), &ty::TyInfer(ty::IntVar(_))) => {
@@ -3445,25 +3470,56 @@ fn check_expr_with_unifier<'a, 'tcx, F>(fcx: &FnCtxt<'a, 'tcx>,
         fcx.write_ty(id, fcx.infcx().next_diverging_ty_var());
       }
       hir::ExprAssign(ref lhs, ref rhs) => {
-        check_expr_with_lvalue_pref(fcx, &**lhs, PreferMutLvalue);
+        if let hir::ExprIndex(ref base, ref idx) = lhs.node {
+            let lvalue_pref = PreferMutLvalue;
+
+            check_expr_with_lvalue_pref(fcx, &**base, lvalue_pref);
+            check_expr(fcx, &**idx);
+
+            let base_t = structurally_resolved_type(fcx, lhs.span, fcx.expr_ty(&**base));
+            let idx_t = fcx.expr_ty(&**idx);
+
+            lookup_indexing(fcx, expr, base, base_t, idx_t, lvalue_pref, true)
+                .map(|(index_ty, rhs_ty)| {
+                    demand::eqtype(fcx, idx.span, index_ty, idx_t);
+                    check_expr_coercable_to_type(fcx, &**rhs, rhs_ty);
+                    fcx.write_nil(lhs.id);
+                    fcx.write_nil(id);
+
+                    if !tcx.sess.features.borrow().indexed_assignments {
+                        tcx.sess.span_err(
+                            expr.span,
+                            "overloaded indexed assignments are not stable");
+                        fileline_help!(
+                            tcx.sess,
+                            expr.span,
+                            "add `#![feature(indexed_assignments)]` to the crate features to \
+                             enable");
+                    }
+                })
+        } else {
+            None
+        }.unwrap_or_else(|| {
+            check_expr_with_lvalue_pref(fcx, &**lhs, PreferMutLvalue);
 
-        let tcx = fcx.tcx();
-        if !tcx.expr_is_lval(&**lhs) {
-            span_err!(tcx.sess, expr.span, E0070,
-                "invalid left-hand side expression");
-        }
+            let tcx = fcx.tcx();
+            if !tcx.expr_is_lval(&**lhs) {
+                span_err!(tcx.sess, expr.span, E0070,
+                    "invalid left-hand side expression");
+            }
 
-        let lhs_ty = fcx.expr_ty(&**lhs);
-        check_expr_coercable_to_type(fcx, &**rhs, lhs_ty);
-        let rhs_ty = fcx.expr_ty(&**rhs);
+            let lhs_ty = fcx.expr_ty(&**lhs);
+            check_expr_coercable_to_type(fcx, &**rhs, lhs_ty);
+            let rhs_ty = fcx.expr_ty(&**rhs);
 
-        fcx.require_expr_have_sized_type(&**lhs, traits::AssignmentLhsSized);
+            fcx.require_expr_have_sized_type(&**lhs, traits::AssignmentLhsSized);
 
-        if lhs_ty.references_error() || rhs_ty.references_error() {
-            fcx.write_error(id);
-        } else {
-            fcx.write_nil(id);
-        }
+            if lhs_ty.references_error() || rhs_ty.references_error() {
+                fcx.write_error(id);
+            } else {
+                fcx.write_nil(id);
+            }
+        })
       }
       hir::ExprIf(ref cond, ref then_blk, ref opt_else_expr) => {
         check_then_else(fcx, &**cond, &**then_blk, opt_else_expr.as_ref().map(|e| &**e),
@@ -3666,7 +3722,7 @@ fn check_expr_with_unifier<'a, 'tcx, F>(fcx: &FnCtxt<'a, 'tcx>,
               fcx.write_ty(id, idx_t);
           } else {
               let base_t = structurally_resolved_type(fcx, expr.span, base_t);
-              match lookup_indexing(fcx, expr, base, base_t, idx_t, lvalue_pref) {
+              match lookup_indexing(fcx, expr, base, base_t, idx_t, lvalue_pref, false) {
                   Some((index_ty, element_ty)) => {
                       let idx_expr_ty = fcx.expr_ty(idx);
                       demand::eqtype(fcx, expr.span, index_ty, idx_expr_ty);