Add rty::FunctionAbi and attach it to rty::FunctionType

Author: coord-e

src/refine/template.rs

@@ -210,6 +210,11 @@ impl<'tcx> TypeBuilder<'tcx> {
         registry: &'a mut R,
         sig: mir_ty::FnSig<'tcx>,
     ) -> FunctionTemplateTypeBuilder<'tcx, 'a, R> {
+        let abi = match sig.abi {
+            rustc_target::spec::abi::Abi::Rust => rty::FunctionAbi::Rust,
+            rustc_target::spec::abi::Abi::RustCall => rty::FunctionAbi::RustCall,
+            _ => unimplemented!("unsupported function ABI: {:?}", sig.abi),
+        };
         FunctionTemplateTypeBuilder {
             inner: self.clone(),
             registry,
@@ -225,6 +230,7 @@ impl<'tcx> TypeBuilder<'tcx> {
             param_rtys: Default::default(),
             param_refinement: None,
             ret_rty: None,
+            abi,
         }
     }
 }
@@ -344,6 +350,7 @@ where
             param_rtys: Default::default(),
             param_refinement: None,
             ret_rty: None,
+            abi: Default::default(),
         }
         .build();
         BasicBlockType { ty, locals }
@@ -359,6 +366,7 @@ pub struct FunctionTemplateTypeBuilder<'tcx, 'a, R> {
     param_refinement: Option<rty::Refinement<rty::FunctionParamIdx>>,
     param_rtys: HashMap<rty::FunctionParamIdx, rty::RefinedType<rty::FunctionParamIdx>>,
     ret_rty: Option<rty::RefinedType<rty::FunctionParamIdx>>,
+    abi: rty::FunctionAbi,
 }
 
 impl<'tcx, 'a, R> FunctionTemplateTypeBuilder<'tcx, 'a, R> {
@@ -478,6 +486,6 @@ where
                 .with_scope(&builder)
                 .build_refined(self.ret_ty)
         });
-        rty::FunctionType::new(param_rtys, ret_rty)
+        rty::FunctionType::new(param_rtys, ret_rty).with_abi(self.abi)
     }
 }

src/rty.rs

@@ -83,6 +83,36 @@ where
     }
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
+pub enum FunctionAbi {
+    #[default]
+    Rust,
+    RustCall,
+}
+
+impl std::fmt::Display for FunctionAbi {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        f.write_str(self.name())
+    }
+}
+
+impl FunctionAbi {
+    pub fn name(&self) -> &'static str {
+        match self {
+            FunctionAbi::Rust => "rust",
+            FunctionAbi::RustCall => "rust-call",
+        }
+    }
+
+    pub fn is_rust(&self) -> bool {
+        matches!(self, FunctionAbi::Rust)
+    }
+
+    pub fn is_rust_call(&self) -> bool {
+        matches!(self, FunctionAbi::RustCall)
+    }
+}
+
 /// A function type.
 ///
 /// In Thrust, function types are closed. Because of that, function types, thus its parameters and
@@ -92,6 +122,7 @@ where
 pub struct FunctionType {
     pub params: IndexVec<FunctionParamIdx, RefinedType<FunctionParamIdx>>,
     pub ret: Box<RefinedType<FunctionParamIdx>>,
+    pub abi: FunctionAbi,
 }
 
 impl<'a, 'b, D> Pretty<'a, D, termcolor::ColorSpec> for &'b FunctionType
@@ -100,15 +131,25 @@ where
     D::Doc: Clone,
 {
     fn pretty(self, allocator: &'a D) -> pretty::DocBuilder<'a, D, termcolor::ColorSpec> {
+        let abi = match self.abi {
+            FunctionAbi::Rust => allocator.nil(),
+            abi => allocator
+                .text("extern")
+                .append(allocator.space())
+                .append(allocator.as_string(abi))
+                .append(allocator.space()),
+        };
         let separator = allocator.text(",").append(allocator.line());
-        allocator
-            .intersperse(self.params.iter().map(|ty| ty.pretty(allocator)), separator)
-            .parens()
-            .append(allocator.space())
-            .append(allocator.text("→"))
-            .append(allocator.line())
-            .append(self.ret.pretty(allocator))
-            .group()
+        abi.append(
+            allocator
+                .intersperse(self.params.iter().map(|ty| ty.pretty(allocator)), separator)
+                .parens(),
+        )
+        .append(allocator.space())
+        .append(allocator.text("→"))
+        .append(allocator.line())
+        .append(self.ret.pretty(allocator))
+        .group()
     }
 }
 
@@ -120,9 +161,15 @@ impl FunctionType {
         FunctionType {
             params,
             ret: Box::new(ret),
+            abi: FunctionAbi::Rust,
         }
     }
 
+    pub fn with_abi(mut self, abi: FunctionAbi) -> Self {
+        self.abi = abi;
+        self
+    }
+
     /// Because function types are always closed in Thrust, we can convert this into
     /// [`Type<Closed>`].
     pub fn into_closed_ty(self) -> Type<Closed> {