Merge pull request #1699 from rtfeldman/effect-after

Fix closures only monomorphizing once

cli/src/repl/eval.rs

@@ -353,6 +353,13 @@ fn jit_to_ast_help<'a>(
         | Layout::RecursivePointer => {
             todo!("add support for rendering recursive tag unions in the REPL")
         }
+        Layout::LambdaSet(lambda_set) => jit_to_ast_help(
+            env,
+            lib,
+            main_fn_name,
+            &lambda_set.runtime_representation(),
+            content,
+        ),
     }
 }
 

cli/tests/cli_run.rs

@@ -242,7 +242,7 @@ mod cli_run {
             filename: "Main.roc",
             executable_filename: "effect-example",
             stdin: &["hi there!"],
-            expected_ending: "hi there!\n",
+            expected_ending: "hi there!\nIt is known\n",
             use_valgrind: true,
         },
         // tea:"tea" => Example {

compiler/gen_llvm/src/llvm/build.rs

@@ -1164,8 +1164,16 @@ pub fn build_exp_expr<'a, 'ctx, 'env>(
         StructAtIndex {
             index, structure, ..
         } => {
+            let (value, layout) = load_symbol_and_layout(scope, structure);
+
+            let layout = if let Layout::LambdaSet(lambda_set) = layout {
+                lambda_set.runtime_representation()
+            } else {
+                *layout
+            };
+
             // extract field from a record
-            match load_symbol_and_layout(scope, structure) {
+            match (value, layout) {
                 (StructValue(argument), Layout::Struct(fields)) => {
                     debug_assert!(!fields.is_empty());
                     env.builder
@@ -4100,6 +4108,11 @@ fn roc_function_call<'a, 'ctx, 'env>(
 ) -> RocFunctionCall<'ctx> {
     use crate::llvm::bitcode::{build_inc_n_wrapper, build_transform_caller};
 
+    let closure_data_layout = match closure_data_layout {
+        Layout::LambdaSet(lambda_set) => lambda_set.runtime_representation(),
+        _ => panic!("closure argument is not a lambda set!"),
+    };
+
     let closure_data_ptr = env
         .builder
         .build_alloca(closure_data.get_type(), "closure_data_ptr");
@@ -4503,8 +4516,13 @@ fn run_higher_order_low_level<'a, 'ctx, 'env>(
 
                     let argument_layouts = &[**element_layout, **element_layout];
 
+                    let closure_data_layout = match closure_layout {
+                        Layout::LambdaSet(lambda_set) => lambda_set.runtime_representation(),
+                        _ => panic!("closure argument is not a lambda set!"),
+                    };
+
                     let compare_wrapper =
-                        build_compare_wrapper(env, function, *closure_layout, element_layout)
+                        build_compare_wrapper(env, function, closure_data_layout, element_layout)
                             .as_global_value()
                             .as_pointer_value();
 

compiler/gen_llvm/src/llvm/build_hash.rs

@@ -59,6 +59,15 @@ fn build_hash_layout<'a, 'ctx, 'env>(
             val.into_struct_value(),
         ),
 
+        Layout::LambdaSet(lambda_set) => build_hash_layout(
+            env,
+            layout_ids,
+            seed,
+            val,
+            &lambda_set.runtime_representation(),
+            when_recursive,
+        ),
+
         Layout::Union(union_layout) => {
             build_hash_tag(env, layout_ids, layout, union_layout, seed, val)
         }

compiler/gen_llvm/src/llvm/compare.rs

@@ -156,6 +156,8 @@ fn build_eq<'a, 'ctx, 'env>(
             rhs_val.into_struct_value(),
         ),
 
+        Layout::LambdaSet(_) => unreachable!("cannot compare closures"),
+
         Layout::Union(union_layout) => build_tag_eq(
             env,
             layout_ids,
@@ -336,6 +338,7 @@ fn build_neq<'a, 'ctx, 'env>(
         Layout::RecursivePointer => {
             unreachable!("recursion pointers should never be compared directly")
         }
+        Layout::LambdaSet(_) => unreachable!("cannot compare closure"),
     }
 }
 

compiler/gen_llvm/src/llvm/convert.rs

@@ -27,6 +27,7 @@ pub fn basic_type_from_layout<'a, 'ctx, 'env>(
 
     match layout {
         Struct(sorted_fields) => basic_type_from_record(env, sorted_fields),
+        LambdaSet(lambda_set) => basic_type_from_layout(env, &lambda_set.runtime_representation()),
         Union(union_layout) => {
             use UnionLayout::*;
 

compiler/gen_llvm/src/llvm/refcounting.rs

@@ -626,6 +626,14 @@ fn modify_refcount_layout_build_function<'a, 'ctx, 'env>(
                 Some(function)
             }
         },
+        LambdaSet(lambda_set) => modify_refcount_layout_build_function(
+            env,
+            parent,
+            layout_ids,
+            mode,
+            when_recursive,
+            &lambda_set.runtime_representation(),
+        ),
     }
 }
 

compiler/mono/src/alias_analysis.rs

@@ -1191,6 +1191,11 @@ fn layout_spec_help(
     match layout {
         Builtin(builtin) => builtin_spec(builder, builtin, when_recursive),
         Struct(fields) => build_recursive_tuple_type(builder, fields, when_recursive),
+        LambdaSet(lambda_set) => layout_spec_help(
+            builder,
+            &lambda_set.runtime_representation(),
+            when_recursive,
+        ),
         Union(union_layout) => {
             let variant_types = build_variant_types(builder, union_layout)?;
 

compiler/mono/src/ir.rs

@@ -1843,7 +1843,7 @@ fn generate_runtime_error_function<'a>(
                 args.push((*arg, env.unique_symbol()));
             }
 
-            args.push((lambda_set.runtime_representation(), Symbol::ARG_CLOSURE));
+            args.push((Layout::LambdaSet(lambda_set), Symbol::ARG_CLOSURE));
 
             (args.into_bump_slice(), *ret_layout)
         }
@@ -1935,12 +1935,11 @@ fn specialize_external<'a>(
                     );
 
                     let body = let_empty_struct(unit, env.arena.alloc(body));
+                    let lambda_set_layout = Layout::LambdaSet(lambda_set);
 
                     let proc = Proc {
                         name,
-                        args: env
-                            .arena
-                            .alloc([(lambda_set.runtime_representation(), Symbol::ARG_CLOSURE)]),
+                        args: env.arena.alloc([(lambda_set_layout, Symbol::ARG_CLOSURE)]),
                         body,
                         closure_data_layout: None,
                         ret_layout: *return_layout,
@@ -1951,7 +1950,7 @@ fn specialize_external<'a>(
 
                     let top_level = ProcLayout::new(
                         env.arena,
-                        env.arena.alloc([lambda_set.runtime_representation()]),
+                        env.arena.alloc([lambda_set_layout]),
                         *return_layout,
                     );
 
@@ -1999,7 +1998,7 @@ fn specialize_external<'a>(
             env.subs.rollback_to(snapshot);
 
             let closure_data_layout = match opt_closure_layout {
-                Some(closure_layout) => closure_layout.runtime_representation(),
+                Some(lambda_set) => Layout::LambdaSet(lambda_set),
                 None => Layout::Struct(&[]),
             };
 
@@ -2149,7 +2148,7 @@ fn specialize_external<'a>(
             env.subs.rollback_to(snapshot);
 
             let closure_data_layout = match opt_closure_layout {
-                Some(closure_layout) => Some(closure_layout.runtime_representation()),
+                Some(lambda_set) => Some(Layout::LambdaSet(lambda_set)),
                 None => None,
             };
 
@@ -2260,7 +2259,7 @@ fn build_specialized_proc<'a>(
         Some(lambda_set) if pattern_symbols.last() == Some(&Symbol::ARG_CLOSURE) => {
             // here we define the lifted (now top-level) f function. Its final argument is `Symbol::ARG_CLOSURE`,
             // it stores the closure structure (just an integer in this case)
-            proc_args.push((lambda_set.runtime_representation(), Symbol::ARG_CLOSURE));
+            proc_args.push((Layout::LambdaSet(lambda_set), Symbol::ARG_CLOSURE));
 
             debug_assert_eq!(
                 pattern_layouts_len + 1,
@@ -2297,7 +2296,7 @@ fn build_specialized_proc<'a>(
                 }
                 Ordering::Greater => {
                     if pattern_symbols.is_empty() {
-                        let ret_layout = lambda_set.runtime_representation();
+                        let ret_layout = Layout::LambdaSet(lambda_set);
                         Ok(FunctionPointerBody {
                             closure: None,
                             ret_layout,
@@ -2482,7 +2481,7 @@ where
     let raw = if procs.module_thunks.contains(&proc_name) {
         match raw {
             RawFunctionLayout::Function(_, lambda_set, _) => {
-                RawFunctionLayout::ZeroArgumentThunk(lambda_set.runtime_representation())
+                RawFunctionLayout::ZeroArgumentThunk(Layout::LambdaSet(lambda_set))
             }
             _ => raw,
         }
@@ -2656,6 +2655,7 @@ macro_rules! match_on_closure_argument {
         let ret_layout = top_level.result;
 
 
+
         match closure_data_layout {
             RawFunctionLayout::Function(_, lambda_set, _) =>  {
                 lowlevel_match_on_lambda_set(
@@ -4131,6 +4131,8 @@ fn construct_closure_data<'a>(
     assigned: Symbol,
     hole: &'a Stmt<'a>,
 ) -> Stmt<'a> {
+    let lambda_set_layout = Layout::LambdaSet(lambda_set);
+
     match lambda_set.layout_for_member(name) {
         ClosureRepresentation::Union {
             tag_id,
@@ -4162,12 +4164,7 @@ fn construct_closure_data<'a>(
                 arguments: symbols,
             };
 
-            Stmt::Let(
-                assigned,
-                expr,
-                lambda_set.runtime_representation(),
-                env.arena.alloc(hole),
-            )
+            Stmt::Let(assigned, expr, lambda_set_layout, env.arena.alloc(hole))
         }
         ClosureRepresentation::AlphabeticOrderStruct(field_layouts) => {
             debug_assert_eq!(field_layouts.len(), symbols.len());
@@ -4198,7 +4195,7 @@ fn construct_closure_data<'a>(
 
             let expr = Expr::Struct(symbols);
 
-            Stmt::Let(assigned, expr, lambda_set.runtime_representation(), hole)
+            Stmt::Let(assigned, expr, lambda_set_layout, hole)
         }
         ClosureRepresentation::Other(Layout::Builtin(Builtin::Int1)) => {
             debug_assert_eq!(symbols.len(), 0);
@@ -4207,7 +4204,7 @@ fn construct_closure_data<'a>(
             let tag_id = name != lambda_set.set[0].0;
             let expr = Expr::Literal(Literal::Bool(tag_id));
 
-            Stmt::Let(assigned, expr, lambda_set.runtime_representation(), hole)
+            Stmt::Let(assigned, expr, lambda_set_layout, hole)
         }
         ClosureRepresentation::Other(Layout::Builtin(Builtin::Int8)) => {
             debug_assert_eq!(symbols.len(), 0);
@@ -4216,7 +4213,7 @@ fn construct_closure_data<'a>(
             let tag_id = lambda_set.set.iter().position(|(s, _)| *s == name).unwrap() as u8;
             let expr = Expr::Literal(Literal::Byte(tag_id));
 
-            Stmt::Let(assigned, expr, lambda_set.runtime_representation(), hole)
+            Stmt::Let(assigned, expr, lambda_set_layout, hole)
         }
         _ => unreachable!(),
     }
@@ -6060,7 +6057,7 @@ fn reuse_function_symbol<'a>(
                         let layout = match raw {
                             RawFunctionLayout::ZeroArgumentThunk(layout) => layout,
                             RawFunctionLayout::Function(_, lambda_set, _) => {
-                                lambda_set.runtime_representation()
+                                Layout::LambdaSet(lambda_set)
                             }
                         };
 
@@ -6158,7 +6155,7 @@ fn reuse_function_symbol<'a>(
                         // TODO suspicious
                         // let layout = Layout::Closure(argument_layouts, lambda_set, ret_layout);
                         // panic!("suspicious");
-                        let layout = lambda_set.runtime_representation();
+                        let layout = Layout::LambdaSet(lambda_set);
                         let top_level = ProcLayout::new(env.arena, &[], layout);
                         procs.insert_passed_by_name(
                             env,
@@ -6348,7 +6345,7 @@ fn call_by_name<'a>(
                         procs,
                         fn_var,
                         proc_name,
-                        env.arena.alloc(lambda_set.runtime_representation()),
+                        env.arena.alloc(Layout::LambdaSet(lambda_set)),
                         layout_cache,
                         assigned,
                         hole,
@@ -6392,7 +6389,7 @@ fn call_by_name<'a>(
                         procs,
                         fn_var,
                         proc_name,
-                        env.arena.alloc(lambda_set.runtime_representation()),
+                        env.arena.alloc(Layout::LambdaSet(lambda_set)),
                         layout_cache,
                         closure_data_symbol,
                         env.arena.alloc(result),
@@ -6522,7 +6519,7 @@ fn call_by_name_help<'a>(
             force_thunk(
                 env,
                 proc_name,
-                lambda_set.runtime_representation(),
+                Layout::LambdaSet(lambda_set),
                 assigned,
                 hole,
             )
@@ -6836,13 +6833,7 @@ fn call_specialized_proc<'a>(
                     arguments: field_symbols,
                 };
 
-                build_call(
-                    env,
-                    call,
-                    assigned,
-                    lambda_set.runtime_representation(),
-                    hole,
-                )
+                build_call(env, call, assigned, Layout::LambdaSet(lambda_set), hole)
             }
             RawFunctionLayout::ZeroArgumentThunk(_) => {
                 unreachable!()
@@ -7909,14 +7900,16 @@ fn match_on_lambda_set<'a>(
     assigned: Symbol,
     hole: &'a Stmt<'a>,
 ) -> Stmt<'a> {
+    let lambda_set_layout = Layout::LambdaSet(lambda_set);
+
     match lambda_set.runtime_representation() {
         Layout::Union(union_layout) => {
             let closure_tag_id_symbol = env.unique_symbol();
 
             let result = union_lambda_set_to_switch(
                 env,
-                lambda_set.set,
-                lambda_set.runtime_representation(),
+                lambda_set,
+                lambda_set_layout,
                 closure_tag_id_symbol,
                 union_layout.tag_id_layout(),
                 closure_data_symbol,
@@ -7946,6 +7939,7 @@ fn match_on_lambda_set<'a>(
             union_lambda_set_branch_help(
                 env,
                 function_symbol,
+                lambda_set,
                 closure_data_symbol,
                 Layout::Struct(fields),
                 argument_symbols,
@@ -7962,7 +7956,7 @@ fn match_on_lambda_set<'a>(
                 env,
                 lambda_set.set,
                 closure_tag_id_symbol,
-                Layout::Builtin(Builtin::Int1),
+                lambda_set_layout,
                 closure_data_symbol,
                 argument_symbols,
                 argument_layouts,
@@ -7978,7 +7972,7 @@ fn match_on_lambda_set<'a>(
                 env,
                 lambda_set.set,
                 closure_tag_id_symbol,
-                Layout::Builtin(Builtin::Int8),
+                lambda_set_layout,
                 closure_data_symbol,
                 argument_symbols,
                 argument_layouts,
@@ -7994,7 +7988,7 @@ fn match_on_lambda_set<'a>(
 #[allow(clippy::too_many_arguments)]
 fn union_lambda_set_to_switch<'a>(
     env: &mut Env<'a, '_>,
-    lambda_set: &'a [(Symbol, &'a [Layout<'a>])],
+    lambda_set: LambdaSet<'a>,
     closure_layout: Layout<'a>,
     closure_tag_id_symbol: Symbol,
     closure_tag_id_layout: Layout<'a>,
@@ -8005,7 +7999,7 @@ fn union_lambda_set_to_switch<'a>(
     assigned: Symbol,
     hole: &'a Stmt<'a>,
 ) -> Stmt<'a> {
-    if lambda_set.is_empty() {
+    if lambda_set.set.is_empty() {
         // NOTE this can happen if there is a type error somewhere. Since the lambda set is empty,
         // there is really nothing we can do here. We generate a runtime error here which allows
         // code gen to proceed. We then assume that we hit another (more descriptive) error before
@@ -8017,11 +8011,12 @@ fn union_lambda_set_to_switch<'a>(
 
     let join_point_id = JoinPointId(env.unique_symbol());
 
-    let mut branches = Vec::with_capacity_in(lambda_set.len(), env.arena);
+    let mut branches = Vec::with_capacity_in(lambda_set.set.len(), env.arena);
 
-    for (i, (function_symbol, _)) in lambda_set.iter().enumerate() {
+    for (i, (function_symbol, _)) in lambda_set.set.iter().enumerate() {
         let stmt = union_lambda_set_branch(
             env,
+            lambda_set,
             join_point_id,
             *function_symbol,
             closure_data_symbol,
@@ -8064,6 +8059,7 @@ fn union_lambda_set_to_switch<'a>(
 #[allow(clippy::too_many_arguments)]
 fn union_lambda_set_branch<'a>(
     env: &mut Env<'a, '_>,
+    lambda_set: LambdaSet<'a>,
     join_point_id: JoinPointId,
     function_symbol: Symbol,
     closure_data_symbol: Symbol,
@@ -8079,6 +8075,7 @@ fn union_lambda_set_branch<'a>(
     union_lambda_set_branch_help(
         env,
         function_symbol,
+        lambda_set,
         closure_data_symbol,
         closure_data_layout,
         argument_symbols_slice,
@@ -8093,6 +8090,7 @@ fn union_lambda_set_branch<'a>(
 fn union_lambda_set_branch_help<'a>(
     env: &mut Env<'a, '_>,
     function_symbol: Symbol,
+    lambda_set: LambdaSet<'a>,
     closure_data_symbol: Symbol,
     closure_data_layout: Layout<'a>,
     argument_symbols_slice: &'a [Symbol],
@@ -8110,7 +8108,7 @@ fn union_lambda_set_branch_help<'a>(
             let mut argument_layouts =
                 Vec::with_capacity_in(argument_layouts_slice.len() + 1, env.arena);
             argument_layouts.extend(argument_layouts_slice);
-            argument_layouts.push(closure_data_layout);
+            argument_layouts.push(Layout::LambdaSet(lambda_set));
 
             // extend symbols with the symbol of the closure environment
             let mut argument_symbols =

compiler/mono/src/layout.rs

@@ -189,6 +189,7 @@ pub enum Layout<'a> {
     /// this is important for closures that capture zero-sized values
     Struct(&'a [Layout<'a>]),
     Union(UnionLayout<'a>),
+    LambdaSet(LambdaSet<'a>),
     RecursivePointer,
 }
 
@@ -531,7 +532,7 @@ impl<'a> LambdaSet<'a> {
                 _ => {
                     let mut arguments = Vec::with_capacity_in(argument_layouts.len() + 1, arena);
                     arguments.extend(argument_layouts);
-                    arguments.push(self.runtime_representation());
+                    arguments.push(Layout::LambdaSet(*self));
 
                     arguments.into_bump_slice()
                 }
@@ -826,6 +827,7 @@ impl<'a> Layout<'a> {
                     }
                 }
             }
+            LambdaSet(lambda_set) => lambda_set.runtime_representation().safe_to_memcpy(),
             RecursivePointer => {
                 // We cannot memcpy pointers, because then we would have the same pointer in multiple places!
                 false
@@ -890,6 +892,9 @@ impl<'a> Layout<'a> {
                     | NonNullableUnwrapped(_) => pointer_size,
                 }
             }
+            LambdaSet(lambda_set) => lambda_set
+                .runtime_representation()
+                .stack_size_without_alignment(pointer_size),
             RecursivePointer => pointer_size,
         }
     }
@@ -919,6 +924,9 @@ impl<'a> Layout<'a> {
                     | NonNullableUnwrapped(_) => pointer_size,
                 }
             }
+            Layout::LambdaSet(lambda_set) => lambda_set
+                .runtime_representation()
+                .alignment_bytes(pointer_size),
             Layout::Builtin(builtin) => builtin.alignment_bytes(pointer_size),
             Layout::RecursivePointer => pointer_size,
         }
@@ -929,6 +937,9 @@ impl<'a> Layout<'a> {
             Layout::Builtin(builtin) => builtin.allocation_alignment_bytes(pointer_size),
             Layout::Struct(_) => unreachable!("not heap-allocated"),
             Layout::Union(union_layout) => union_layout.allocation_alignment_bytes(pointer_size),
+            Layout::LambdaSet(lambda_set) => lambda_set
+                .runtime_representation()
+                .allocation_alignment_bytes(pointer_size),
             Layout::RecursivePointer => unreachable!("should be looked up to get an actual layout"),
         }
     }
@@ -979,6 +990,7 @@ impl<'a> Layout<'a> {
                     | NonNullableUnwrapped(_) => true,
                 }
             }
+            LambdaSet(lambda_set) => lambda_set.runtime_representation().contains_refcounted(),
             RecursivePointer => true,
         }
     }
@@ -1002,6 +1014,7 @@ impl<'a> Layout<'a> {
                     .append(alloc.text("}"))
             }
             Union(union_layout) => union_layout.to_doc(alloc, parens),
+            LambdaSet(lambda_set) => lambda_set.runtime_representation().to_doc(alloc, parens),
             RecursivePointer => alloc.text("*self"),
         }
     }
@@ -1360,7 +1373,7 @@ fn layout_from_flat_type<'a>(
         Func(_, closure_var, _) => {
             let lambda_set = LambdaSet::from_var(env.arena, env.subs, closure_var, env.ptr_bytes)?;
 
-            Ok(lambda_set.runtime_representation())
+            Ok(Layout::LambdaSet(lambda_set))
         }
         Record(fields, ext_var) => {
             // extract any values from the ext_var

compiler/test_gen/src/gen_primitives.rs

@@ -2618,7 +2618,8 @@ fn lambda_set_struct_byte() {
                     r = Red
 
                     p1 = (\u -> r == u)
-                    oneOfResult = List.map [p1, p1] (\p -> p Green)
+                    foobarbaz = (\p -> p Green)
+                    oneOfResult = List.map [p1, p1] foobarbaz
 
                     when oneOfResult is
                         _ -> 32

examples/effect/Main.roc

@@ -5,4 +5,7 @@ app "effect-example"
 
 main : Effect.Effect {}
 main =
-    Effect.after Effect.getLine \lineThisThing -> Effect.putLine lineThisThing
+    Effect.after (Effect.getLine) \line ->
+        Effect.after (Effect.putLine "You entered: \(line)") \{} ->
+            Effect.after (Effect.putLine "It is known") \{} ->
+                Effect.always {}