Add LambdaCaptureKind to allow specifying move/copy/borrow treatment …

…for free variables. Tweak simplifier to *not* remove lambda values (even though they are now ) to avoid losing some errors/warnings.
aptos-labs · Oct 18, 2024 · e24bf4e · e24bf4e
1 parent 3f1f0a6
commit e24bf4e
Show file tree

Hide file tree

Showing 18 changed files with 570 additions and 358 deletions.
diff --git a/third_party/move/move-compiler-v2/src/bytecode_generator.rs b/third_party/move/move-compiler-v2/src/bytecode_generator.rs
@@ -480,7 +480,7 @@ impl<'env> Generator<'env> {
                 self.emit_with(*id, |attr| Bytecode::SpecBlock(attr, spec));
             },
             // TODO(LAMBDA)
-            ExpData::Lambda(id, _, _, _) => self.error(
+            ExpData::Lambda(id, _, _, _, _) => self.error(
                 *id,
                 "Function-typed values not yet supported except as parameters to calls to inline functions",
             ),
@@ -814,7 +814,7 @@ impl<'env> Generator<'env> {
             Operation::NoOp => {}, // do nothing
 
             // TODO(LAMBDA)
-            Operation::Closure(..) => self.error(
+            Operation::Closure => self.error(
                 id,
                 "Function-typed values not yet supported except as parameters to calls to inline functions",
             ),
@@ -1331,12 +1331,9 @@ impl<'env> Generator<'env> {
         };
         self.gen_borrow_field_operation(id, borrow_dest, str, fields, oper_temp);
         if need_read_ref {
-            self.emit_call(
-                id,
-                vec![target],
-                BytecodeOperation::ReadRef,
-                vec![borrow_dest],
-            )
+            self.emit_call(id, vec![target], BytecodeOperation::ReadRef, vec![
+                borrow_dest,
+            ])
         }
     }
 
@@ -1524,13 +1521,10 @@ enum MatchMode {
 impl MatchMode {
     /// Whether this match is in probing mode.
     fn is_probing(&self) -> bool {
-        matches!(
-            self,
-            MatchMode::Refutable {
-                probing_vars: Some(_),
-                ..
-            }
-        )
+        matches!(self, MatchMode::Refutable {
+            probing_vars: Some(_),
+            ..
+        })
     }
 
     /// Whether a variable appearing in the pattern should be bound to a temporary.
@@ -1658,12 +1652,9 @@ impl<'env> Generator<'env> {
                         ReferenceKind::Immutable,
                         Box::new(value_ty.clone()),
                     ));
-                    self.emit_call(
-                        id,
-                        vec![value_ref],
-                        BytecodeOperation::BorrowLoc,
-                        vec![value],
-                    );
+                    self.emit_call(id, vec![value_ref], BytecodeOperation::BorrowLoc, vec![
+                        value,
+                    ]);
                     needs_probing = true;
                     value_ref
                 }
@@ -1785,11 +1776,10 @@ impl<'env> Generator<'env> {
                             ),
                         );
                         return Some(
-                            ExpData::Call(
-                                id,
-                                Operation::Deref,
-                                vec![ExpData::LocalVar(new_id, var).into_exp()],
+                            ExpData::Call(id, Operation::Deref, vec![ExpData::LocalVar(
+                                new_id, var,
                             )
+                            .into_exp()])
                             .into_exp(),
                         );
                     }

diff --git a/third_party/move/move-compiler-v2/src/env_pipeline/ast_simplifier.rs b/third_party/move/move-compiler-v2/src/env_pipeline/ast_simplifier.rs
@@ -359,7 +359,7 @@ fn find_possibly_modified_vars(
                     _ => {},
                 }
             },
-            Lambda(node_id, pat, _, _) => {
+            Lambda(node_id, pat, _, _, _) => {
                 // Define a new scope for bound vars, and turn off `modifying` within.
                 match pos {
                     VisitorPosition::Pre => {
@@ -978,7 +978,8 @@ impl<'env> ExpRewriterFunctions for SimplifierRewriter<'env> {
                     let ability_set = self
                         .env()
                         .type_abilities(&ty, self.func_env.get_type_parameters_ref());
-                    ability_set.has_ability(Ability::Drop)
+                    // Don't drop a function-valued expression so we don't lose errors.
+                    !ty.has_function() && ability_set.has_ability(Ability::Drop)
                 } else {
                     // We're missing type info, be conservative
                     false

diff --git a/third_party/move/move-compiler-v2/src/env_pipeline/flow_insensitive_checkers.rs b/third_party/move/move-compiler-v2/src/env_pipeline/flow_insensitive_checkers.rs
@@ -147,7 +147,7 @@ impl<'env, 'params> SymbolVisitor<'env, 'params> {
                 Pre | Post | BeforeBody | MidMutate | BeforeThen | BeforeElse
                 | PreSequenceValue => {},
             },
-            Lambda(_, pat, _, _) => {
+            Lambda(_, pat, _, _, _) => {
                 match position {
                     Pre => self.seen_uses.enter_scope(),
                     Post => {

diff --git a/third_party/move/move-compiler-v2/src/env_pipeline/function_checker.rs b/third_party/move/move-compiler-v2/src/env_pipeline/function_checker.rs
@@ -270,7 +270,7 @@ fn check_privileged_operations_on_structs(env: &GlobalEnv, fun_env: &FunctionEnv
                 },
                 ExpData::Assign(_, pat, _)
                 | ExpData::Block(_, pat, _, _)
-                | ExpData::Lambda(_, pat, _, _) => {
+                | ExpData::Lambda(_, pat, _, _, _) => {
                     pat.visit_pre_post(&mut |_, pat| {
                         if let Pattern::Struct(id, str, _, _) = pat {
                             let module_id = str.module_id;

diff --git a/third_party/move/move-compiler-v2/src/env_pipeline/inliner.rs b/third_party/move/move-compiler-v2/src/env_pipeline/inliner.rs
@@ -1161,7 +1161,7 @@ impl<'env, 'rewriter> ExpRewriterFunctions for InlinedRewriter<'env, 'rewriter>
         };
         let call_loc = self.env.get_node_loc(id);
         if let Some(lambda_target) = optional_lambda_target {
-            if let ExpData::Lambda(_, pat, body, _) = lambda_target.as_ref() {
+            if let ExpData::Lambda(_, pat, body, _, _) = lambda_target.as_ref() {
                 let args_vec: Vec<Exp> = args.to_vec();
                 Some(InlinedRewriter::construct_inlined_call_expression(
                     self.env,

diff --git a/third_party/move/move-compiler-v2/src/env_pipeline/lambda_lifter.rs b/third_party/move/move-compiler-v2/src/env_pipeline/lambda_lifter.rs
@@ -40,10 +40,9 @@
 //! ```
 
 use itertools::Itertools;
-use move_binary_format::file_format::AbilitySet;
-use move_binary_format::file_format::Visibility;
+use move_binary_format::file_format::{AbilitySet, Visibility};
 use move_model::{
-    ast::{Exp, ExpData, Operation, Pattern, TempIndex},
+    ast::{Exp, ExpData, LambdaCaptureKind, Operation, Pattern, TempIndex},
     exp_rewriter::{ExpRewriter, ExpRewriterFunctions, RewriteTarget},
     model::{FunId, FunctionEnv, GlobalEnv, Loc, NodeId, Parameter, TypeParameter},
     symbol::Symbol,
@@ -260,13 +259,10 @@ impl<'a> ExpRewriterFunctions for LambdaLifter<'a> {
 
     fn rewrite_assign(&mut self, _node_id: NodeId, lhs: &Pattern, _rhs: &Exp) -> Option<Exp> {
         for (node_id, name) in lhs.vars() {
-            self.free_locals.insert(
-                name,
-                VarInfo {
-                    node_id,
-                    modified: true,
-                },
-            );
+            self.free_locals.insert(name, VarInfo {
+                node_id,
+                modified: true,
+            });
         }
         None
     }
@@ -275,22 +271,16 @@ impl<'a> ExpRewriterFunctions for LambdaLifter<'a> {
         if matches!(oper, Operation::Borrow(ReferenceKind::Mutable)) {
             match args[0].as_ref() {
                 ExpData::LocalVar(node_id, name) => {
-                    self.free_locals.insert(
-                        *name,
-                        VarInfo {
-                            node_id: *node_id,
-                            modified: true,
-                        },
-                    );
+                    self.free_locals.insert(*name, VarInfo {
+                        node_id: *node_id,
+                        modified: true,
+                    });
                 },
                 ExpData::Temporary(node_id, param) => {
-                    self.free_params.insert(
-                        *param,
-                        VarInfo {
-                            node_id: *node_id,
-                            modified: true,
-                        },
-                    );
+                    self.free_params.insert(*param, VarInfo {
+                        node_id: *node_id,
+                        modified: true,
+                    });
                 },
                 _ => {},
             }
@@ -303,6 +293,7 @@ impl<'a> ExpRewriterFunctions for LambdaLifter<'a> {
         id: NodeId,
         pat: &Pattern,
         body: &Exp,
+        capture_kind: LambdaCaptureKind,
         _abilities: AbilitySet, // TODO(LAMBDA): do something with this
     ) -> Option<Exp> {
         if self.exempted_lambdas.contains(&id) {
@@ -315,6 +306,21 @@ impl<'a> ExpRewriterFunctions for LambdaLifter<'a> {
         // parameter indices in the lambda context to indices in the lifted
         // functions (courtesy of #12317)
         let mut param_index_mapping = BTreeMap::new();
+        match capture_kind {
+            LambdaCaptureKind::Default => {
+                // OK.
+            },
+            LambdaCaptureKind::Move | LambdaCaptureKind::Copy | LambdaCaptureKind::Borrow => {
+                let loc = env.get_node_loc(id);
+                env.error(
+                    &loc,
+                    &format!(
+                        "Lambda function `{}` of free variables not yet supported.", // TODO(LAMBDA)
+                        capture_kind
+                    ),
+                );
+            },
+        };
         for (used_param_count, (param, var_info)) in
             mem::take(&mut self.free_params).into_iter().enumerate()
         {

diff --git a/third_party/move/move-compiler/src/attr_derivation/mod.rs b/third_party/move/move-compiler/src/attr_derivation/mod.rs
@@ -218,7 +218,10 @@ pub(crate) fn new_full_name(
 
 /// Helper to create a call exp.
 pub(crate) fn new_call_exp(loc: Loc, fun: NameAccessChain, args: Vec<Exp>) -> Exp {
-    sp(loc, Exp_::Call(fun, CallKind::Regular, None, sp(loc, args)))
+    sp(
+        loc,
+        Exp_::Call(fun, CallKind::Regular, None, sp(loc, args), false),
+    )
 }
 
 pub(crate) fn new_borrow_exp(loc: Loc, arg: Exp) -> Exp {

diff --git a/third_party/move/move-compiler/src/expansion/ast.rs b/third_party/move/move-compiler/src/expansion/ast.rs
@@ -6,8 +6,8 @@ use crate::{
     expansion::translate::is_valid_struct_constant_or_schema_name,
     parser::ast::{
         self as P, Ability, Ability_, BinOp, CallKind, ConstantName, Field, FunctionName, Label,
-        ModuleName, QuantKind, SpecApplyPattern, StructName, UnaryOp, UseDecl, Var, VariantName,
-        ENTRY_MODIFIER,
+        LambdaCaptureKind, ModuleName, QuantKind, SpecApplyPattern, StructName, UnaryOp, UseDecl,
+        Var, VariantName, ENTRY_MODIFIER,
     },
     shared::{
         ast_debug::*,
@@ -497,8 +497,18 @@ pub enum Exp_ {
     Copy(Var),
 
     Name(ModuleAccess, Option<Vec<Type>>),
-    Call(ModuleAccess, CallKind, Option<Vec<Type>>, Spanned<Vec<Exp>>),
-    ExpCall(Box<Exp>, Spanned<Vec<Exp>>),
+    Call(
+        ModuleAccess,
+        CallKind,
+        Option<Vec<Type>>,
+        Spanned<Vec<Exp>>,
+        bool, // ends in ".."
+    ),
+    ExpCall(
+        Box<Exp>,
+        Spanned<Vec<Exp>>,
+        bool, // ends in ".."
+    ),
     Pack(ModuleAccess, Option<Vec<Type>>, Fields<Exp>),
     Vector(Loc, Option<Vec<Type>>, Spanned<Vec<Exp>>),
 
@@ -507,7 +517,7 @@ pub enum Exp_ {
     While(Option<Label>, Box<Exp>, Box<Exp>),
     Loop(Option<Label>, Box<Exp>),
     Block(Sequence),
-    Lambda(TypedLValueList, Box<Exp>, AbilitySet),
+    Lambda(TypedLValueList, Box<Exp>, LambdaCaptureKind, AbilitySet),
     Quant(
         QuantKind,
         LValueWithRangeList,
@@ -1611,7 +1621,7 @@ impl AstDebug for Exp_ {
                     w.write(">");
                 }
             },
-            E::Call(ma, kind, tys_opt, sp!(_, rhs)) => {
+            E::Call(ma, kind, tys_opt, sp!(_, rhs), ends_in_dotdot) => {
                 ma.ast_debug(w);
                 w.write(kind.to_string());
                 if let Some(ss) = tys_opt {
@@ -1621,12 +1631,18 @@ impl AstDebug for Exp_ {
                 }
                 w.write("(");
                 w.comma(rhs, |w, e| e.ast_debug(w));
+                if ends_in_dotdot {
+                    w.write("..");
+                }
                 w.write(")");
             },
-            E::ExpCall(fexp, sp!(_, rhs)) => {
+            E::ExpCall(fexp, sp!(_, rhs), ends_in_dotdot) => {
                 fexp.ast_debug(w);
                 w.write("(");
                 w.comma(rhs, |w, e| e.ast_debug(w));
+                if ends_in_dotdot {
+                    w.write("..");
+                }
                 w.write(")");
             },
             E::Pack(ma, tys_opt, fields) => {
@@ -1694,7 +1710,10 @@ impl AstDebug for Exp_ {
                 }
             },
             E::Block(seq) => w.block(|w| seq.ast_debug(w)),
-            E::Lambda(sp!(_, bs), e, abilities) => {
+            E::Lambda(sp!(_, bs), e, capture_kind, abilities) => {
+                if *capture_kind != LambdaCaptureKind::Default {
+                    w.write(format!(" {}", capture_kind));
+                }
                 w.write("|");
                 bs.ast_debug(w);
                 w.write("|");

diff --git a/third_party/move/move-compiler/src/expansion/dependency_ordering.rs b/third_party/move/move-compiler/src/expansion/dependency_ordering.rs
@@ -446,12 +446,12 @@ fn exp(context: &mut Context, sp!(_loc, e_): &E::Exp) {
             module_access(context, ma);
             types_opt(context, tys_opt)
         },
-        E::Call(ma, _is_macro, tys_opt, sp!(_, args_)) => {
+        E::Call(ma, _is_macro, tys_opt, sp!(_, args_), _ends_in_dotdot) => {
             module_access(context, ma);
             types_opt(context, tys_opt);
             args_.iter().for_each(|e| exp(context, e))
         },
-        E::ExpCall(fexp, sp!(_, args_)) => {
+        E::ExpCall(fexp, sp!(_, args_), _ends_in_dotdot) => {
             exp(context, fexp);
             args_.iter().for_each(|e| exp(context, e))
         },
@@ -516,7 +516,7 @@ fn exp(context: &mut Context, sp!(_loc, e_): &E::Exp) {
             tys.iter().for_each(|ty| type_(context, ty))
         },
 
-        E::Lambda(ll, e, _abilities) => {
+        E::Lambda(ll, e, _capture_kind, _abilities) => {
             use crate::expansion::ast::TypedLValue_;
             let mapped = ll.value.iter().map(|sp!(_, TypedLValue_(lv, _opt_ty))| lv);
             lvalues(context, mapped);