Extract mutable references (without drop/reassignment)

stainless_extraction/src/expr/block.rs

@@ -126,10 +126,7 @@ impl<'a, 'l, 'tcx> BodyExtractor<'a, 'l, 'tcx> {
                 let body_expr = self.extract_block_(stmts, acc_exprs, acc_specs, final_expr);
 
                 // wrap that body expression into the Let
-
-                // This is safe, because we don't support mutable references.
                 let init = self.extract_aliasable_expr(init);
-
                 let last_expr = if vd.is_mutable() {
                   f.LetVar(vd, init, body_expr).into()
                 } else {

stainless_extraction/src/expr/field.rs

@@ -25,6 +25,14 @@ impl<'a, 'l, 'tcx> BodyExtractor<'a, 'l, 'tcx> {
         ),
       },
 
+      ExprKind::Deref { arg } => {
+        let arg = self.extract_expr(arg);
+        self
+          .factory()
+          .FieldAssignment(arg, self.synth().mut_ref_value_id(), value)
+          .into()
+      }
+
       e => self.unsupported_expr(
         lhs.span,
         format!("Cannot extract assignment to kind {:?}", e),

stainless_extraction/src/expr/mod.rs

@@ -1,7 +1,7 @@
 use std::convert::TryFrom;
 
 use literal::Literal;
-use rustc_middle::mir::BorrowKind;
+use rustc_middle::mir::{BorrowKind, Mutability};
 use rustc_middle::ty::{subst::SubstsRef, Ty, TyKind};
 use rustc_mir_build::thir::{BlockSafety, Expr, ExprKind, FruInfo, Stmt, StmtKind};
 
@@ -66,7 +66,16 @@ impl<'a, 'l, 'tcx> BodyExtractor<'a, 'l, 'tcx> {
       ExprKind::Use { source } => self.extract_expr(source),
       ExprKind::NeverToAny { source } => self.extract_expr(source),
 
-      ExprKind::Deref { arg } => self.extract_expr(arg),
+      ExprKind::Deref { arg } => match arg.ty.ref_mutability() {
+        Some(Mutability::Mut) => {
+          let arg = self.extract_expr(arg);
+          self
+            .factory()
+            .ADTSelector(arg, self.synth().mut_ref_value_id())
+            .into()
+        }
+        _ => self.extract_expr(arg),
+      },
 
       // Borrow an immutable and aliasable value (i.e. the meaning of
       // BorrowKind::Shared). Handle this safe case with erasure of the
@@ -77,6 +86,15 @@ impl<'a, 'l, 'tcx> BodyExtractor<'a, 'l, 'tcx> {
         arg,
       } => self.extract_aliasable_expr(arg),
 
+      ExprKind::Borrow {
+        borrow_kind: BorrowKind::Mut { .. },
+        arg,
+      } => {
+        let tpe = self.base.extract_ty(arg.ty, &self.txtcx, arg.span);
+        let arg = self.extract_expr(arg);
+        self.synth().mut_ref(tpe, arg)
+      }
+
       ExprKind::Assign { lhs, rhs } => self.extract_assignment(lhs, rhs),
 
       ExprKind::Return { value } => self.extract_return(*value),
@@ -358,17 +376,23 @@ impl<'a, 'l, 'tcx> BodyExtractor<'a, 'l, 'tcx> {
   }
 
   /// Inserts a fresh copy around the expression, if the expression can contain
-  /// mutable types.
-  /// Currently, this is safe everywhere because:
-  ///  - we either work on an immutable reference aka shared borrow aka `&T`
+  /// mutable types. And if the expression IS NOT a mutable reference
+  /// This is safe because:
+  ///  - we don't do anything for mutable references.
+  ///  - Otherwise, we either work on an immutable reference aka shared borrow aka `&T`
   ///  - or we own the the value, meaning it's moved when we return/alias it somewhere
   ///    and the compiler has made sure that it's not used afterwards.
   ///
-  /// TODO: Check that this is still safe (at call sites) when mutable references are introduced.
   fn extract_aliasable_expr(&mut self, expr: &'a Expr<'a, 'tcx>) -> st::Expr<'l> {
     let e = self.extract_expr(expr);
-    let tpe = self.base.extract_ty(expr.ty, &self.txtcx, expr.span);
-    self.base.fresh_copy_if_needed(e, tpe)
+
+    // If this is a mutable reference, we DON'T freshCopy
+    if let Some(Mutability::Mut) = expr.ty.ref_mutability() {
+      e
+    } else {
+      let tpe = self.base.extract_ty(expr.ty, &self.txtcx, expr.span);
+      self.base.fresh_copy_if_needed(e, tpe)
+    }
   }
 
   fn unsupported_expr<M: Into<String>>(&mut self, span: Span, msg: M) -> st::Expr<'l> {

stainless_extraction/src/krate.rs

@@ -359,10 +359,6 @@ impl<'l, 'tcx> BaseExtractor<'l, 'tcx> {
     let (carrier_flags, mut flags_by_symbol) = self.extract_flags(hir_id);
     let mut flags = carrier_flags.to_stainless(f);
 
-    // As long as we only have local mutation, every function viewed from the
-    // outside is pure.
-    flags.push(f.IsPure().into());
-
     // Add flag specifying that this function is a method of its class (if there's a class)
     flags.extend(
       class_def
@@ -396,6 +392,9 @@ impl<'l, 'tcx> BaseExtractor<'l, 'tcx> {
         (bxtor.dcx.params().to_vec(), bxtor.return_tpe(), body_expr)
       });
 
+    // TODO: check whether the function has any '&mut' params, otherwise mark as 'pure'.
+    flags.push(f.IsPure().into());
+
     self.report_unused_flags(hir_id, &flags_by_symbol);
 
     // Wrap it all up in a Stainless function

stainless_extraction/src/lib.rs

@@ -99,6 +99,7 @@ struct Extraction<'l> {
   generics: HashMap<StainlessSymId<'l>, Generics<'l>>,
   functions: HashMap<StainlessSymId<'l>, &'l st::FunDef<'l>>,
   classes: HashMap<StainlessSymId<'l>, &'l st::ClassDef<'l>>,
+  mut_ref_id: Option<StainlessSymId<'l>>,
 }
 
 impl<'l> Extraction<'l> {
@@ -112,6 +113,7 @@ impl<'l> Extraction<'l> {
       generics: Default::default(),
       functions: Default::default(),
       classes: Default::default(),
+      mut_ref_id: Default::default(),
     }
   }
 

stainless_extraction/src/synth/mod.rs

@@ -0,0 +1,33 @@
+use super::*;
+
+mod mut_ref;
+mod std_option;
+
+/// This is just a proxy, a mutable reference to the base extractor. The idea of
+/// `Synth` is just to scope some methods under a common namespace, therefore
+/// the struct mainly provides a namespace/scope.
+pub struct Synth<'a, 'l, 'tcx> {
+  base: &'a mut BaseExtractor<'l, 'tcx>,
+}
+
+impl<'l, 'tcx> BaseExtractor<'l, 'tcx> {
+  /// Returns a `Synth` which provides access to all methods that synthesise
+  /// items/trees.
+  pub fn synth<'a>(&'a mut self) -> Synth<'a, 'l, 'tcx> {
+    Synth { base: self }
+  }
+}
+
+impl<'b, 'l, 'tcx> BodyExtractor<'b, 'l, 'tcx> {
+  /// Returns a `Synth` which provides access to all methods that synthesise
+  /// items/trees.
+  pub fn synth<'a>(&'a mut self) -> Synth<'a, 'l, 'tcx> {
+    self.base.synth()
+  }
+}
+
+impl<'l> Synth<'_, 'l, '_> {
+  fn factory(&self) -> &'l st::Factory {
+    self.base.factory()
+  }
+}

stainless_extraction/src/synth/mut_ref.rs

@@ -0,0 +1,53 @@
+use super::*;
+use stainless_data::ast::ADTSort;
+
+impl<'a, 'l, 'tcx> Synth<'a, 'l, 'tcx> {
+  pub fn mut_ref(&mut self, tpe: st::Type<'l>, arg: st::Expr<'l>) -> st::Expr<'l> {
+    self
+      .factory()
+      .ADT(self.mut_ref_id(), vec![tpe], vec![arg])
+      .into()
+  }
+
+  pub fn mut_ref_type(&mut self, tpe: st::Type<'l>) -> st::Type<'l> {
+    self.factory().ADTType(self.mut_ref_id(), vec![tpe]).into()
+  }
+
+  pub fn mut_ref_value_id(&mut self) -> StainlessSymId<'l> {
+    self.get_or_create_adt().constructors[0].fields[0].v.id
+  }
+
+  fn mut_ref_id(&mut self) -> StainlessSymId<'l> {
+    self.get_or_create_adt().id
+  }
+
+  fn get_or_create_adt(&mut self) -> &'l ADTSort<'l> {
+    self
+      .base
+      .with_extraction(|xt| xt.mut_ref_id.and_then(|id| xt.adts.get(id).copied()))
+      .unwrap_or_else(|| {
+        let f = self.factory();
+
+        let adt_id = self.base.fresh_id("MutRef".into());
+        self
+          .base
+          .with_extraction_mut(|xt| xt.mut_ref_id = Some(adt_id));
+
+        let field_id = self.base.fresh_id("t".into());
+        let tparam = &*f.TypeParameter(self.base.fresh_id("T".into()), vec![f.IsMutable().into()]);
+
+        let adt = f.ADTSort(
+          adt_id,
+          vec![f.TypeParameterDef(tparam)],
+          vec![f.ADTConstructor(
+            adt_id,
+            adt_id,
+            vec![f.ValDef(f.Variable(field_id, tparam.into(), vec![f.IsVar().into()]))],
+          )],
+          vec![f.Synthetic().into()],
+        );
+        self.base.add_adt(adt);
+        adt
+      })
+  }
+}

stainless_extraction/src/synth/std_option.rs

@@ -0,0 +1,53 @@
+use super::*;
+
+/// Synthetisation of std::option::Option trees in Stainless AST.
+impl<'a, 'l, 'tcx> Synth<'a, 'l, 'tcx> {
+  pub fn std_option_type(&mut self, tpe: st::Type<'l>) -> st::Type<'l> {
+    self
+      .factory()
+      .ADTType(self.option_adt().id, vec![tpe])
+      .into()
+  }
+
+  pub fn std_option_none(&mut self, tpe: st::Type<'l>) -> st::Expr<'l> {
+    self.factory().ADT(self.none_id(), vec![tpe], vec![]).into()
+  }
+
+  pub fn std_option_some(&mut self, val: st::Expr<'l>, tpe: st::Type<'l>) -> st::Expr<'l> {
+    self
+      .factory()
+      .ADT(self.some_id(), vec![tpe], vec![val])
+      .into()
+  }
+
+  pub fn std_option_some_type(&mut self, tpe: st::Type<'l>) -> st::Type<'l> {
+    self.factory().ADTType(self.some_id(), vec![tpe]).into()
+  }
+
+  pub fn std_option_some_value(&mut self, some: st::Expr<'l>) -> st::Expr<'l> {
+    self
+      .factory()
+      .ClassSelector(some, self.some_value_id())
+      .into()
+  }
+
+  fn option_adt(&mut self) -> &'l st::ADTSort<'l> {
+    let def_id = self
+      .base
+      .std_items
+      .item_to_def(StdItem::CrateItem(CrateItem::OptionType));
+    self.base.get_or_extract_adt(def_id)
+  }
+
+  fn none_id(&mut self) -> StainlessSymId<'l> {
+    self.option_adt().constructors[0].id
+  }
+
+  fn some_id(&mut self) -> StainlessSymId<'l> {
+    self.option_adt().constructors[1].id
+  }
+
+  fn some_value_id(&mut self) -> StainlessSymId<'l> {
+    self.option_adt().constructors[1].fields[0].v.id
+  }
+}

stainless_extraction/src/ty.rs

@@ -147,6 +147,11 @@ impl<'l, 'tcx> BaseExtractor<'l, 'tcx> {
       // Immutable references
       TyKind::Ref(_, ty, Mutability::Not) => self.extract_ty(ty, txtcx, span),
 
+      TyKind::Ref(_, ty, Mutability::Mut) => {
+        let arg_ty = self.extract_ty(ty, txtcx, span);
+        self.synth().mut_ref_type(arg_ty)
+      }
+
       TyKind::Param(param_ty) => txtcx
         .index_to_tparam
         .get(&param_ty.index)