Typecheck and compile opaque wrapping functions

This enables you to write something like

```
A := U8
List.map [1, 2, 3] @A
```

which will be compiled as if it was `List.map [1, 2, 3] \x -> @A x`.

Closes #3499

crates/compiler/can/src/expr.rs

@@ -220,6 +220,9 @@ pub enum Expr {
         lambda_set_variables: Vec<LambdaSet>,
     },
 
+    // Opaque as a function, e.g. @Id as a shorthand for \x -> @Id x
+    OpaqueWrapFunction(OpaqueWrapFunctionData),
+
     // Test
     Expect {
         loc_condition: Box<Loc<Expr>>,
@@ -269,6 +272,9 @@ impl Expr {
                 args_count: 0,
             },
             &Self::OpaqueRef { name, .. } => Category::OpaqueWrap(name),
+            &Self::OpaqueWrapFunction(OpaqueWrapFunctionData { opaque_name, .. }) => {
+                Category::OpaqueWrap(opaque_name)
+            }
             Self::Expect { .. } => Category::Expect,
 
             // these nodes place no constraints on the expression's type
@@ -380,6 +386,76 @@ impl AccessorData {
     }
 }
 
+/// An opaque wrapper like `@Foo`, which is equivalent to `\p -> @Foo p`
+/// These are desugared to closures, but we distinguish them so we can have
+/// better error messages during constraint generation.
+#[derive(Clone, Debug)]
+pub struct OpaqueWrapFunctionData {
+    pub opaque_name: Symbol,
+    pub opaque_var: Variable,
+    // The following fields help link the concrete opaque type; see
+    // `Expr::OpaqueRef` for more info on how they're used.
+    pub specialized_def_type: Type,
+    pub type_arguments: Vec<OptAbleVar>,
+    pub lambda_set_variables: Vec<LambdaSet>,
+
+    pub function_name: Symbol,
+    pub function_var: Variable,
+    pub argument_var: Variable,
+    pub closure_var: Variable,
+}
+
+impl OpaqueWrapFunctionData {
+    pub fn to_closure_data(self, argument_symbol: Symbol) -> ClosureData {
+        let OpaqueWrapFunctionData {
+            opaque_name,
+            opaque_var,
+            specialized_def_type,
+            type_arguments,
+            lambda_set_variables,
+            function_name,
+            function_var,
+            argument_var,
+            closure_var,
+        } = self;
+
+        // IDEA: convert
+        //
+        // @Foo
+        //
+        // into
+        //
+        // (\p -> @Foo p)
+        let body = Expr::OpaqueRef {
+            opaque_var,
+            name: opaque_name,
+            argument: Box::new((argument_var, Loc::at_zero(Expr::Var(argument_symbol)))),
+            specialized_def_type: Box::new(specialized_def_type),
+            type_arguments,
+            lambda_set_variables,
+        };
+
+        let loc_body = Loc::at_zero(body);
+
+        let arguments = vec![(
+            argument_var,
+            AnnotatedMark::known_exhaustive(),
+            Loc::at_zero(Pattern::Identifier(argument_symbol)),
+        )];
+
+        ClosureData {
+            function_type: function_var,
+            closure_type: closure_var,
+            return_type: opaque_var,
+            name: function_name,
+            captured_symbols: vec![],
+            recursive: Recursive::NotRecursive,
+            arguments,
+            loc_body: Box::new(loc_body),
+        }
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct Field {
     pub var: Variable,
@@ -835,15 +911,41 @@ pub fn canonicalize_expr<'a>(
                 Output::default(),
             )
         }
-        ast::Expr::OpaqueRef(opaque_ref) => {
+        ast::Expr::OpaqueRef(name) => {
             // If we're here, the opaque reference is definitely not wrapping an argument - wrapped
             // arguments are handled in the Apply branch.
-            let problem = roc_problem::can::RuntimeError::OpaqueNotApplied(Loc::at(
-                region,
-                (*opaque_ref).into(),
-            ));
-            env.problem(Problem::RuntimeError(problem.clone()));
-            (RuntimeError(problem), Output::default())
+            // Treat this as a function \payload -> @Opaque payload
+            match scope.lookup_opaque_ref(name, region) {
+                Err(runtime_error) => {
+                    env.problem(Problem::RuntimeError(runtime_error.clone()));
+                    (RuntimeError(runtime_error), Output::default())
+                }
+                Ok((name, opaque_def)) => {
+                    let mut output = Output::default();
+                    output.references.insert_type_lookup(name);
+
+                    let (type_arguments, lambda_set_variables, specialized_def_type) =
+                        freshen_opaque_def(var_store, opaque_def);
+
+                    let fn_symbol = scope.gen_unique_symbol();
+
+                    (
+                        OpaqueWrapFunction(OpaqueWrapFunctionData {
+                            opaque_name: name,
+                            opaque_var: var_store.fresh(),
+                            specialized_def_type,
+                            type_arguments,
+                            lambda_set_variables,
+
+                            function_name: fn_symbol,
+                            function_var: var_store.fresh(),
+                            argument_var: var_store.fresh(),
+                            closure_var: var_store.fresh(),
+                        }),
+                        output,
+                    )
+                }
+            }
         }
         ast::Expr::Expect(condition, continuation) => {
             let mut output = Output::default();
@@ -1420,7 +1522,8 @@ pub fn inline_calls(var_store: &mut VarStore, scope: &mut Scope, expr: Expr) ->
         | other @ AbilityMember(..)
         | other @ RunLowLevel { .. }
         | other @ TypedHole { .. }
-        | other @ ForeignCall { .. } => other,
+        | other @ ForeignCall { .. }
+        | other @ OpaqueWrapFunction(_) => other,
 
         List {
             elem_var,
@@ -2387,7 +2490,8 @@ fn get_lookup_symbols(expr: &Expr, var_store: &mut VarStore) -> Vec<(Symbol, Var
             | Expr::SingleQuote(_)
             | Expr::EmptyRecord
             | Expr::TypedHole(_)
-            | Expr::RuntimeError(_) => {}
+            | Expr::RuntimeError(_)
+            | Expr::OpaqueWrapFunction(_) => {}
         }
     }
 

crates/compiler/can/src/module.rs

@@ -970,5 +970,6 @@ fn fix_values_captured_in_closure_expr(
             let (_, loc_arg) = &mut **argument;
             fix_values_captured_in_closure_expr(&mut loc_arg.value, no_capture_symbols);
         }
+        OpaqueWrapFunction(_) => {}
     }
 }

crates/compiler/can/src/traverse.rs

@@ -7,7 +7,10 @@ use roc_types::subs::Variable;
 use crate::{
     abilities::AbilitiesStore,
     def::{Annotation, Declaration, Def},
-    expr::{self, AccessorData, AnnotatedMark, ClosureData, Declarations, Expr, Field},
+    expr::{
+        self, AccessorData, AnnotatedMark, ClosureData, Declarations, Expr, Field,
+        OpaqueWrapFunctionData,
+    },
     pattern::{DestructType, Pattern, RecordDestruct},
 };
 
@@ -220,6 +223,7 @@ pub fn walk_expr<V: Visitor>(visitor: &mut V, expr: &Expr, var: Variable) {
             ext_var: _,
         } => visitor.visit_expr(&loc_expr.value, loc_expr.region, *field_var),
         Expr::Accessor(AccessorData { .. }) => { /* terminal */ }
+        Expr::OpaqueWrapFunction(OpaqueWrapFunctionData { .. }) => { /* terminal */ }
         Expr::Update {
             record_var: _,
             ext_var: _,

crates/compiler/constrain/src/expr.rs

@@ -13,7 +13,7 @@ use roc_can::expected::PExpected;
 use roc_can::expr::Expr::{self, *};
 use roc_can::expr::{
     AccessorData, AnnotatedMark, ClosureData, DeclarationTag, Declarations, DestructureDef, Field,
-    FunctionDef, WhenBranch,
+    FunctionDef, OpaqueWrapFunctionData, WhenBranch,
 };
 use roc_can::pattern::Pattern;
 use roc_can::traverse::symbols_introduced_from_pattern;
@@ -1114,6 +1114,100 @@ pub fn constrain_expr(
 
             constraints.exists_many(vars, [arg_con, opaque_con, link_type_variables_con])
         }
+        OpaqueWrapFunction(OpaqueWrapFunctionData {
+            opaque_name,
+            opaque_var,
+            specialized_def_type,
+            type_arguments,
+            lambda_set_variables,
+            function_name,
+            function_var,
+            argument_var,
+            closure_var,
+        }) => {
+            let argument_type = Type::Variable(*argument_var);
+
+            let opaque_type = Type::Alias {
+                symbol: *opaque_name,
+                type_arguments: type_arguments
+                    .iter()
+                    .map(|v| OptAbleType {
+                        typ: Type::Variable(v.var),
+                        opt_ability: v.opt_ability,
+                    })
+                    .collect(),
+                lambda_set_variables: lambda_set_variables.clone(),
+                actual: Box::new(argument_type.clone()),
+                kind: AliasKind::Opaque,
+            };
+
+            // Tie the opaque type to the opaque_var
+            let opaque_con = constraints.equal_types_var(
+                *opaque_var,
+                Expected::NoExpectation(opaque_type),
+                Category::OpaqueWrap(*opaque_name),
+                region,
+            );
+
+            // Tie the type of the value wrapped by the opaque to the opaque's type variables.
+            let link_type_variables_con = constraints.equal_types(
+                argument_type.clone(),
+                Expected::NoExpectation((*specialized_def_type).clone()),
+                Category::OpaqueArg,
+                region,
+            );
+
+            let lambda_set = Type::ClosureTag {
+                name: *function_name,
+                captures: vec![],
+                ambient_function: *function_var,
+            };
+
+            let closure_type = Type::Variable(*closure_var);
+
+            let opaque_type = Type::Variable(*opaque_var);
+
+            let function_type = Type::Function(
+                vec![argument_type],
+                Box::new(closure_type),
+                Box::new(opaque_type),
+            );
+
+            let cons = [
+                opaque_con,
+                link_type_variables_con,
+                constraints.equal_types_var(
+                    *closure_var,
+                    NoExpectation(lambda_set),
+                    Category::OpaqueWrap(*opaque_name),
+                    region,
+                ),
+                constraints.equal_types_var(
+                    *function_var,
+                    Expected::NoExpectation(function_type),
+                    Category::OpaqueWrap(*opaque_name),
+                    region,
+                ),
+                constraints.equal_types_var(
+                    *function_var,
+                    expected,
+                    Category::OpaqueWrap(*opaque_name),
+                    region,
+                ),
+            ];
+
+            let mut vars = vec![*argument_var, *opaque_var];
+
+            // Also add the fresh variables we created for the type argument and lambda sets
+            vars.extend(type_arguments.iter().map(|v| v.var));
+            vars.extend(lambda_set_variables.iter().map(|v| {
+                v.0.expect_variable("all lambda sets should be fresh variables here")
+            }));
+
+            vars.extend([*function_var, *closure_var]);
+
+            constraints.exists_many(vars, cons)
+        }
 
         RunLowLevel { args, ret_var, op } => {
             // This is a modified version of what we do for function calls.

crates/compiler/mono/src/copy.rs

@@ -1,6 +1,6 @@
 use roc_can::{
     def::Def,
-    expr::{AccessorData, ClosureData, Expr, Field, WhenBranch},
+    expr::{AccessorData, ClosureData, Expr, Field, OpaqueWrapFunctionData, WhenBranch},
 };
 use roc_module::ident::{Lowercase, TagName};
 use roc_types::{
@@ -543,6 +543,29 @@ fn deep_copy_type_vars_into_expr_help<C: CopyEnv>(
                 lambda_set_variables: lambda_set_variables.clone(),
             },
 
+            OpaqueWrapFunction(OpaqueWrapFunctionData {
+                opaque_name,
+                opaque_var,
+                specialized_def_type,
+                type_arguments,
+                lambda_set_variables,
+                function_name,
+                function_var,
+                argument_var,
+                closure_var,
+            }) => OpaqueWrapFunction(OpaqueWrapFunctionData {
+                opaque_name: *opaque_name,
+                opaque_var: sub!(*opaque_var),
+                function_name: *function_name,
+                function_var: sub!(*function_var),
+                argument_var: sub!(*argument_var),
+                closure_var: sub!(*closure_var),
+                // The following three are only used for constraining
+                specialized_def_type: specialized_def_type.clone(),
+                type_arguments: type_arguments.clone(),
+                lambda_set_variables: lambda_set_variables.clone(),
+            }),
+
             Expect {
                 loc_condition,
                 loc_continuation,

crates/compiler/mono/src/ir.rs

@@ -4387,6 +4387,60 @@ pub fn with_hole<'a>(
             }
         }
 
+        OpaqueWrapFunction(wrap_fn_data) => {
+            let opaque_var = wrap_fn_data.opaque_var;
+            let arg_symbol = env.unique_symbol();
+
+            let ClosureData {
+                name,
+                function_type,
+                arguments,
+                loc_body,
+                ..
+            } = wrap_fn_data.to_closure_data(arg_symbol);
+
+            match procs.insert_anonymous(
+                env,
+                LambdaName::no_niche(name),
+                function_type,
+                arguments,
+                *loc_body,
+                CapturedSymbols::None,
+                opaque_var,
+                layout_cache,
+            ) {
+                Ok(_) => {
+                    let raw_layout = return_on_layout_error!(
+                        env,
+                        layout_cache.raw_from_var(env.arena, function_type, env.subs),
+                        "Expr::OpaqueWrapFunction"
+                    );
+
+                    match raw_layout {
+                        RawFunctionLayout::Function(_, lambda_set, _) => {
+                            let lambda_name =
+                                find_lambda_name(env, layout_cache, lambda_set, name, &[]);
+                            construct_closure_data(
+                                env,
+                                lambda_set,
+                                lambda_name,
+                                &[],
+                                assigned,
+                                hole,
+                            )
+                        }
+                        RawFunctionLayout::ZeroArgumentThunk(_) => {
+                            internal_error!("should not be a thunk!")
+                        }
+                    }
+                }
+
+                Err(_error) => Stmt::RuntimeError(
+                    "TODO convert anonymous function error to a RuntimeError string",
+                ),
+            }
+        }
+
         Update {
             record_var,
             symbol: structure,

crates/compiler/solve/tests/solve_expr.rs

@@ -7341,4 +7341,30 @@ mod solve_expr {
             "Rose I64",
         );
     }
+
+    #[test]
+    fn opaque_wrap_function() {
+        infer_eq_without_problem(
+            indoc!(
+                r#"
+                A := U8
+                List.map [1, 2, 3] @A
+                "#
+            ),
+            "List A",
+        );
+    }
+
+    #[test]
+    fn opaque_wrap_function_with_inferred_arg() {
+        infer_eq_without_problem(
+            indoc!(
+                r#"
+                A a := a
+                List.map [1u8, 2u8, 3u8] @A
+                "#
+            ),
+            "List (A U8)",
+        );
+    }
 }