clean up the fndef air types

source/vir/src/datatype_to_air.rs

@@ -75,6 +75,7 @@ fn uses_ext_equal(ctx: &Ctx, typ: &Typ) -> bool {
         TypX::StrSlice => false,
         TypX::Char => false,
         TypX::Primitive(_, _) => true,
+        TypX::FnDef(..) => false,
     }
 }
 
@@ -672,6 +673,22 @@ pub fn datatypes_and_primitives_to_air(ctx: &Ctx, datatypes: &crate::ast::Dataty
             is_transparent && datatype.x.ext_equal,
         );
     }
+
+    for fun in &ctx.fndef_types {
+        let func = ctx.func_map.get(fun).expect("expected fndef function in pruned crate");
+        let tparams = &func.x.typ_params;
+        let mut args: Vec<air::ast::Typ> = Vec::new();
+        for _ in tparams.iter() {
+            args.extend(crate::def::types().iter().map(|s| str_typ(s)));
+        }
+        let decl_type_id = Arc::new(DeclX::fun_or_const(
+            crate::def::prefix_fndef_type_id(fun),
+            Arc::new(args),
+            str_typ(crate::def::TYPE),
+        ));
+        token_commands.push(Arc::new(CommandX::Global(decl_type_id)));
+    }
+
     let mut commands: Vec<Command> = Vec::new();
     commands.append(&mut opaque_sort_commands);
     commands.push(Arc::new(CommandX::Global(Arc::new(DeclX::Datatypes(Arc::new(

source/vir/src/def.rs

@@ -52,6 +52,7 @@ const PREFIX_PRE_VAR: &str = "pre%";
 const PREFIX_BOX: &str = "Poly%";
 const PREFIX_UNBOX: &str = "%Poly%";
 const PREFIX_TYPE_ID: &str = "TYPE%";
+const PREFIX_FNDEF_TYPE_ID: &str = "FNDEF%";
 const PREFIX_TUPLE_TYPE: &str = "tuple%";
 const PREFIX_CLOSURE_TYPE: &str = "anonymous_closure%";
 const PREFIX_TUPLE_PARAM: &str = "T%";
@@ -62,7 +63,6 @@ const PREFIX_PROJECT: &str = "proj%";
 const PREFIX_PROJECT_DECORATION: &str = "proj%%";
 const PREFIX_TRAIT_BOUND: &str = "tr_bound%";
 const PREFIX_STATIC: &str = "static%";
-const PREFIX_FNDEF_TYPE: &str = "fndef%";
 const SLICE_TYPE: &str = "slice%";
 const ARRAY_TYPE: &str = "array%";
 const PREFIX_SNAPSHOT: &str = "snap%";
@@ -88,7 +88,8 @@ pub const SUFFIX_SNAP_WHILE_END: &str = "_while_end";
 
 pub const CLOSURE_RETURN_VALUE_PREFIX: &str = "%closure_return";
 
-pub const VARIANT_FNDEF_SINGLETON: &str = "fndef_single%";
+pub const FNDEF_TYPE: &str = "fndef";
+pub const FNDEF_SINGLETON: &str = "fndef_singleton";
 
 // List of constant strings that can appear in generated AIR code
 pub const FUEL_ID: &str = "FuelId";
@@ -124,10 +125,12 @@ pub const BOX_INT: &str = "I";
 pub const BOX_BOOL: &str = "B";
 pub const BOX_STRSLICE: &str = "S";
 pub const BOX_CHAR: &str = "C";
+pub const BOX_FNDEF: &str = "F";
 pub const UNBOX_INT: &str = "%I";
 pub const UNBOX_BOOL: &str = "%B";
 pub const UNBOX_STRSLICE: &str = "%S";
 pub const UNBOX_CHAR: &str = "%C";
+pub const UNBOX_FNDEF: &str = "%F";
 pub const TYPE: &str = "Type";
 pub const TYPE_ID_BOOL: &str = "BOOL";
 pub const TYPE_ID_INT: &str = "INT";
@@ -332,6 +335,10 @@ pub fn prefix_type_id(path: &Path) -> Ident {
     Arc::new(PREFIX_TYPE_ID.to_string() + &path_to_string(path))
 }
 
+pub fn prefix_fndef_type_id(fun: &Fun) -> Ident {
+    Arc::new(PREFIX_FNDEF_TYPE_ID.to_string() + &fun_to_string(fun))
+}
+
 pub fn prefix_tuple_type(i: usize) -> Path {
     let ident = Arc::new(format!("{}{}", PREFIX_TUPLE_TYPE, i));
     Arc::new(PathX { krate: None, segments: Arc::new(vec![ident]) })
@@ -382,13 +389,6 @@ pub fn prefix_type_id_fun(i: usize) -> Ident {
     prefix_type_id(&prefix_lambda_type(i))
 }
 
-pub fn prefix_fndef_type(fun: &Fun) -> Path {
-    let krate: Option<Ident> = fun.path.krate.clone();
-    let mut segments: Vec<Ident> = (*fun.path.segments).clone();
-    segments.insert(0, Arc::new(PREFIX_FNDEF_TYPE.to_string()));
-    Arc::new(PathX { krate, segments: Arc::new(segments) })
-}
-
 pub fn prefix_box(ident: &Path) -> Ident {
     Arc::new(PREFIX_BOX.to_string() + &path_to_string(ident))
 }
@@ -397,14 +397,6 @@ pub fn prefix_unbox(ident: &Path) -> Ident {
     Arc::new(PREFIX_UNBOX.to_string() + &path_to_string(ident))
 }
 
-pub fn prefix_fndef_box(fun: &Fun) -> Ident {
-    prefix_box(&prefix_fndef_type(fun))
-}
-
-pub fn prefix_fndef_unbox(fun: &Fun) -> Ident {
-    prefix_unbox(&prefix_fndef_type(fun))
-}
-
 pub fn prefix_fuel_id(ident: &Ident) -> Ident {
     Arc::new(PREFIX_FUEL_ID.to_string() + ident)
 }

source/vir/src/prelude.rs

@@ -53,8 +53,15 @@ pub(crate) fn prelude_nodes(config: PreludeConfig) -> Vec<Node> {
     let Height = str_to_node(T_HEIGHT);
     let box_int = str_to_node(BOX_INT);
     let box_bool = str_to_node(BOX_BOOL);
+    let box_fndef = str_to_node(BOX_FNDEF);
     let unbox_int = str_to_node(UNBOX_INT);
     let unbox_bool = str_to_node(UNBOX_BOOL);
+    let unbox_fndef = str_to_node(UNBOX_FNDEF);
+
+    #[allow(non_snake_case)]
+    let FnDef = str_to_node(FNDEF_TYPE);
+    #[allow(non_snake_case)]
+    let FnDefSingleton = str_to_node(FNDEF_SINGLETON);
 
     let box_strslice = str_to_node(BOX_STRSLICE);
     let unbox_strslice = str_to_node(UNBOX_STRSLICE);
@@ -142,13 +149,18 @@ pub(crate) fn prelude_nodes(config: PreludeConfig) -> Vec<Node> {
         (declare-fun [new_strlit] (Int) [strslice])
         (declare-fun [from_strlit] ([strslice]) Int)
 
+        // FnDef
+        (declare-datatypes (([FnDef] 0)) ((([FnDefSingleton]))))
+
         // Polymorphism
         (declare-sort [Poly] 0)
         (declare-sort [Height] 0)
         (declare-fun [box_int] (Int) [Poly])
         (declare-fun [box_bool] (Bool) [Poly])
+        (declare-fun [box_fndef] ([FnDef]) [Poly])
         (declare-fun [unbox_int] ([Poly]) Int)
         (declare-fun [unbox_bool] ([Poly]) Bool)
+        (declare-fun [unbox_fndef] ([Poly]) [FnDef])
         (declare-fun [box_strslice] ([strslice]) [Poly])
         (declare-fun [unbox_strslice] ([Poly]) [strslice])
         (declare-fun [box_char] ([Char]) [Poly])

source/vir/src/sst_to_air.rs

@@ -133,9 +133,7 @@ pub(crate) fn typ_to_air(ctx: &Ctx, typ: &Typ) -> air::ast::Typ {
             }
         }
         TypX::Decorate(_, t) => typ_to_air(ctx, t),
-        TypX::FnDef(fun, _, _, _) => {
-            ident_typ(&path_to_air_ident(&crate::def::prefix_fndef_type(fun)))
-        }
+        TypX::FnDef(..) => str_typ(crate::def::FNDEF_TYPE),
         TypX::Boxed(_) => str_typ(POLY),
         TypX::TypParam(_) => str_typ(POLY),
         TypX::Primitive(Primitive::Array | Primitive::Slice, _) => match typ_as_mono(typ) {
@@ -314,8 +312,12 @@ pub(crate) fn primitive_id(name: &Primitive, typs: &Typs) -> Expr {
 }
 
 pub(crate) fn fndef_id(fun: &Fun, typs: &Typs) -> Expr {
-    let path = crate::def::prefix_fndef_type(fun);
-    datatype_id(&path, typs)
+    let f_name = crate::def::prefix_fndef_type_id(fun);
+    let mut args: Vec<Expr> = Vec::new();
+    for t in typs.iter() {
+        args.extend(typ_to_ids(t));
+    }
+    air::ast_util::ident_apply_or_var(&f_name, &Arc::new(args))
 }
 
 pub(crate) fn expr_has_type(expr: &Expr, typ: &Expr) -> Expr {
@@ -425,7 +427,7 @@ fn try_box(ctx: &Ctx, expr: Expr, typ: &Typ) -> Option<Expr> {
             }
         }
         TypX::Primitive(_, _) => prefix_typ_as_mono(prefix_box, typ, "primitive type"),
-        TypX::FnDef(fun, _, _, _) => Some(crate::def::prefix_fndef_box(fun)),
+        TypX::FnDef(..) => Some(str_ident(crate::def::BOX_FNDEF)),
         TypX::Decorate(_, t) => return try_box(ctx, expr, t),
         TypX::Boxed(_) => None,
         TypX::TypParam(_) => None,
@@ -455,7 +457,7 @@ fn try_unbox(ctx: &Ctx, expr: Expr, typ: &Typ) -> Option<Expr> {
             }
         }
         TypX::Primitive(_, _) => prefix_typ_as_mono(prefix_unbox, typ, "primitive type"),
-        TypX::FnDef(fun, _, _, _) => Some(crate::def::prefix_fndef_box(fun)),
+        TypX::FnDef(..) => Some(str_ident(crate::def::UNBOX_FNDEF)),
         TypX::Tuple(_) => None,
         TypX::Lambda(typs, _) => Some(prefix_unbox(&prefix_lambda_type(typs.len()))),
         TypX::AnonymousClosure(..) => unimplemented!(),
@@ -772,12 +774,8 @@ pub(crate) fn exp_to_expr(ctx: &Ctx, exp: &Exp, expr_ctxt: &ExprCtxt) -> Result<
                 .collect::<Result<Vec<_>, VirErr>>()?;
             Arc::new(ExprX::Apply(variant, Arc::new(args)))
         }
-        ExpX::ExecFnByName(fun) => {
-            let singleton_ctor = variant_ident(
-                &crate::def::prefix_fndef_type(fun),
-                &Arc::new(crate::def::VARIANT_FNDEF_SINGLETON.to_string()),
-            );
-            Arc::new(ExprX::Apply(singleton_ctor, Arc::new(vec![])))
+        ExpX::ExecFnByName(_fun) => {
+            Arc::new(ExprX::Apply(str_ident(crate::def::FNDEF_SINGLETON), Arc::new(vec![])))
         }
         ExpX::NullaryOpr(crate::ast::NullaryOpr::ConstGeneric(c)) => {
             str_apply(crate::def::CONST_INT, &vec![typ_to_id(c)])