1#![warn(missing_docs)]
58#![warn(clippy::all)]
59#![warn(clippy::pedantic)]
60
61pub mod dyn_tensor;
62pub mod nat;
63pub mod ty_list;
64
65pub use dyn_tensor::{dyn_tensor_of, dyn_tensor_tycon, shape_witness_of, shape_witness_tycon};
66pub use nat::TyNat;
67pub use ty_list::TyList;
68
69use bhc_index::Idx;
70use bhc_intern::Symbol;
71use bhc_span::Span;
72use serde::{Deserialize, Serialize};
73
74#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
79pub struct TyId(u32);
80
81impl Idx for TyId {
82 #[allow(clippy::cast_possible_truncation)]
83 fn new(idx: usize) -> Self {
84 Self(idx as u32)
85 }
86
87 fn index(self) -> usize {
88 self.0 as usize
89 }
90}
91
92#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
98pub struct TyVar {
99 pub id: u32,
101 pub kind: Kind,
103}
104
105impl TyVar {
106 #[must_use]
108 pub fn new(id: u32, kind: Kind) -> Self {
109 Self { id, kind }
110 }
111
112 #[must_use]
114 pub fn new_star(id: u32) -> Self {
115 Self::new(id, Kind::Star)
116 }
117}
118
119#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
127pub enum Kind {
128 Star,
130 Arrow(Box<Kind>, Box<Kind>),
132 Constraint,
134 Var(u32),
136
137 Nat,
146
147 List(Box<Kind>),
154}
155
156impl Kind {
157 #[must_use]
159 pub fn star_to_star() -> Self {
160 Self::Arrow(Box::new(Self::Star), Box::new(Self::Star))
161 }
162
163 #[must_use]
165 pub fn is_star(&self) -> bool {
166 matches!(self, Self::Star)
167 }
168
169 #[must_use]
171 pub fn is_nat(&self) -> bool {
172 matches!(self, Self::Nat)
173 }
174
175 #[must_use]
177 pub fn nat_list() -> Self {
178 Self::List(Box::new(Self::Nat))
179 }
180
181 #[must_use]
183 pub fn tensor_kind() -> Self {
184 Self::Arrow(
186 Box::new(Self::nat_list()),
187 Box::new(Self::star_to_star()),
188 )
189 }
190
191 #[must_use]
193 pub fn is_list(&self) -> bool {
194 matches!(self, Self::List(_))
195 }
196
197 #[must_use]
199 pub fn list_elem_kind(&self) -> Option<&Kind> {
200 match self {
201 Self::List(k) => Some(k),
202 _ => None,
203 }
204 }
205}
206
207#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
212pub enum Ty {
213 Var(TyVar),
215
216 Con(TyCon),
218
219 Prim(PrimTy),
222
223 App(Box<Ty>, Box<Ty>),
225
226 Fun(Box<Ty>, Box<Ty>),
229
230 Tuple(Vec<Ty>),
232
233 List(Box<Ty>),
236
237 Forall(Vec<TyVar>, Box<Ty>),
239
240 Error,
242
243 Nat(TyNat),
249
250 TyList(TyList),
254}
255
256impl Ty {
257 #[must_use]
259 pub fn unit() -> Self {
260 Self::Tuple(Vec::new())
261 }
262
263 #[must_use]
265 pub fn fun(from: Ty, to: Ty) -> Self {
266 Self::Fun(Box::new(from), Box::new(to))
267 }
268
269 #[must_use]
271 pub fn list(elem: Ty) -> Self {
272 Self::List(Box::new(elem))
273 }
274
275 #[must_use]
277 pub fn int_prim() -> Self {
278 Self::Prim(PrimTy::I64)
279 }
280
281 #[must_use]
283 pub fn double_prim() -> Self {
284 Self::Prim(PrimTy::F64)
285 }
286
287 #[must_use]
289 pub fn float_prim() -> Self {
290 Self::Prim(PrimTy::F32)
291 }
292
293 #[must_use]
295 pub fn is_prim(&self) -> bool {
296 matches!(self, Self::Prim(_))
297 }
298
299 #[must_use]
301 pub fn as_prim(&self) -> Option<PrimTy> {
302 match self {
303 Self::Prim(p) => Some(*p),
304 _ => None,
305 }
306 }
307
308 #[must_use]
310 pub fn is_fun(&self) -> bool {
311 matches!(self, Self::Fun(_, _))
312 }
313
314 #[must_use]
316 pub fn is_error(&self) -> bool {
317 matches!(self, Self::Error)
318 }
319
320 #[must_use]
322 pub fn free_vars(&self) -> Vec<TyVar> {
323 let mut vars = Vec::new();
324 self.collect_free_vars(&mut vars);
325 vars
326 }
327
328 fn collect_free_vars(&self, vars: &mut Vec<TyVar>) {
329 match self {
330 Self::Var(v) => {
331 if !vars.contains(v) {
332 vars.push(v.clone());
333 }
334 }
335 Self::Con(_) | Self::Prim(_) | Self::Error => {}
336 Self::App(f, a) | Self::Fun(f, a) => {
337 f.collect_free_vars(vars);
338 a.collect_free_vars(vars);
339 }
340 Self::Tuple(tys) => {
341 for ty in tys {
342 ty.collect_free_vars(vars);
343 }
344 }
345 Self::List(elem) => elem.collect_free_vars(vars),
346 Self::Forall(bound, body) => {
347 let mut body_vars = Vec::new();
348 body.collect_free_vars(&mut body_vars);
349 for v in body_vars {
350 if !bound.contains(&v) && !vars.contains(&v) {
351 vars.push(v);
352 }
353 }
354 }
355 Self::Nat(n) => {
357 for v in n.free_vars() {
358 if !vars.contains(&v) {
359 vars.push(v);
360 }
361 }
362 }
363 Self::TyList(l) => {
364 for v in l.free_vars() {
365 if !vars.contains(&v) {
366 vars.push(v);
367 }
368 }
369 }
370 }
371 }
372
373 #[must_use]
375 pub fn is_ground(&self) -> bool {
376 match self {
377 Self::Var(_) => false,
378 Self::Con(_) | Self::Prim(_) | Self::Error => true,
379 Self::App(f, a) | Self::Fun(f, a) => f.is_ground() && a.is_ground(),
380 Self::Tuple(tys) => tys.iter().all(Ty::is_ground),
381 Self::List(elem) => elem.is_ground(),
382 Self::Forall(_, body) => body.is_ground(),
383 Self::Nat(n) => n.is_ground(),
384 Self::TyList(l) => l.is_ground(),
385 }
386 }
387
388 #[must_use]
390 pub fn nat_lit(n: u64) -> Self {
391 Self::Nat(TyNat::lit(n))
392 }
393
394 #[must_use]
396 pub fn shape(dims: &[u64]) -> Self {
397 Self::TyList(TyList::shape_from_dims(dims))
398 }
399
400 #[must_use]
402 pub fn is_nat(&self) -> bool {
403 matches!(self, Self::Nat(_))
404 }
405
406 #[must_use]
408 pub fn is_ty_list(&self) -> bool {
409 matches!(self, Self::TyList(_))
410 }
411
412 #[must_use]
414 pub fn as_nat(&self) -> Option<&TyNat> {
415 match self {
416 Self::Nat(n) => Some(n),
417 _ => None,
418 }
419 }
420
421 #[must_use]
423 pub fn as_ty_list(&self) -> Option<&TyList> {
424 match self {
425 Self::TyList(l) => Some(l),
426 _ => None,
427 }
428 }
429}
430
431impl std::fmt::Display for Ty {
432 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
433 match self {
434 Self::Var(v) => write!(f, "t{}", v.id),
435 Self::Con(c) => write!(f, "{}", c.name.as_str()),
436 Self::Prim(p) => write!(f, "{p}"),
437 Self::App(fun, arg) => write!(f, "({fun} {arg})"),
438 Self::Fun(from, to) => write!(f, "({from} -> {to})"),
439 Self::Tuple(tys) => {
440 write!(f, "(")?;
441 for (i, ty) in tys.iter().enumerate() {
442 if i > 0 {
443 write!(f, ", ")?;
444 }
445 write!(f, "{ty}")?;
446 }
447 write!(f, ")")
448 }
449 Self::List(elem) => write!(f, "[{elem}]"),
450 Self::Forall(vars, body) => {
451 write!(f, "forall")?;
452 for v in vars {
453 write!(f, " t{}", v.id)?;
454 }
455 write!(f, ". {body}")
456 }
457 Self::Error => write!(f, "<error>"),
458 Self::Nat(n) => write!(f, "{n}"),
459 Self::TyList(l) => write!(f, "{l}"),
460 }
461 }
462}
463
464#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
466pub struct TyCon {
467 pub name: Symbol,
469 pub kind: Kind,
471}
472
473impl TyCon {
474 #[must_use]
476 pub fn new(name: Symbol, kind: Kind) -> Self {
477 Self { name, kind }
478 }
479}
480
481#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
487pub struct Scheme {
488 pub vars: Vec<TyVar>,
490 pub constraints: Vec<Constraint>,
492 pub ty: Ty,
494}
495
496impl Scheme {
497 #[must_use]
499 pub fn mono(ty: Ty) -> Self {
500 Self {
501 vars: Vec::new(),
502 constraints: Vec::new(),
503 ty,
504 }
505 }
506
507 #[must_use]
509 pub fn poly(vars: Vec<TyVar>, ty: Ty) -> Self {
510 Self {
511 vars,
512 constraints: Vec::new(),
513 ty,
514 }
515 }
516
517 #[must_use]
519 pub fn qualified(vars: Vec<TyVar>, constraints: Vec<Constraint>, ty: Ty) -> Self {
520 Self {
521 vars,
522 constraints,
523 ty,
524 }
525 }
526
527 #[must_use]
529 pub fn is_mono(&self) -> bool {
530 self.vars.is_empty()
531 }
532}
533
534#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
536pub struct Constraint {
537 pub class: Symbol,
539 pub args: Vec<Ty>,
541 pub span: Span,
543}
544
545impl Constraint {
546 #[must_use]
548 pub fn new(class: Symbol, ty: Ty, span: Span) -> Self {
549 Self {
550 class,
551 args: vec![ty],
552 span,
553 }
554 }
555
556 #[must_use]
558 pub fn new_multi(class: Symbol, args: Vec<Ty>, span: Span) -> Self {
559 Self { class, args, span }
560 }
561}
562
563#[derive(Clone, Debug, Default, Serialize, Deserialize)]
565pub struct Subst {
566 mapping: rustc_hash::FxHashMap<u32, Ty>,
568}
569
570impl Subst {
571 #[must_use]
573 pub fn new() -> Self {
574 Self::default()
575 }
576
577 pub fn insert(&mut self, var: &TyVar, ty: Ty) {
579 self.mapping.insert(var.id, ty);
580 }
581
582 #[must_use]
584 pub fn get(&self, var: &TyVar) -> Option<&Ty> {
585 self.mapping.get(&var.id)
586 }
587
588 #[must_use]
590 pub fn contains(&self, var: &TyVar) -> bool {
591 self.mapping.contains_key(&var.id)
592 }
593
594 #[must_use]
596 pub fn apply(&self, ty: &Ty) -> Ty {
597 match ty {
598 Ty::Var(v) => self.get(v).cloned().unwrap_or_else(|| ty.clone()),
599 Ty::Con(_) | Ty::Prim(_) => ty.clone(),
600 Ty::App(f, a) => Ty::App(Box::new(self.apply(f)), Box::new(self.apply(a))),
601 Ty::Fun(from, to) => Ty::Fun(Box::new(self.apply(from)), Box::new(self.apply(to))),
602 Ty::Tuple(tys) => Ty::Tuple(tys.iter().map(|t| self.apply(t)).collect()),
603 Ty::List(elem) => Ty::List(Box::new(self.apply(elem))),
604 Ty::Forall(vars, body) => {
605 let mut inner = self.clone();
607 for v in vars {
608 inner.mapping.remove(&v.id);
609 }
610 Ty::Forall(vars.clone(), Box::new(inner.apply(body)))
611 }
612 Ty::Error => Ty::Error,
613 Ty::Nat(n) => Ty::Nat(self.apply_nat(n)),
615 Ty::TyList(l) => Ty::TyList(self.apply_ty_list(l)),
616 }
617 }
618
619 #[must_use]
621 pub fn apply_nat(&self, n: &TyNat) -> TyNat {
622 match n {
623 TyNat::Lit(val) => TyNat::Lit(*val),
624 TyNat::Var(v) => {
625 match self.get(v) {
627 Some(Ty::Nat(replacement)) => replacement.clone(),
628 Some(_) => n.clone(), None => n.clone(),
630 }
631 }
632 TyNat::Add(a, b) => TyNat::add(self.apply_nat(a), self.apply_nat(b)),
633 TyNat::Mul(a, b) => TyNat::mul(self.apply_nat(a), self.apply_nat(b)),
634 }
635 }
636
637 #[must_use]
639 pub fn apply_ty_list(&self, l: &TyList) -> TyList {
640 match l {
641 TyList::Nil => TyList::Nil,
642 TyList::Cons(head, tail) => {
643 TyList::cons(self.apply(head), self.apply_ty_list(tail))
644 }
645 TyList::Var(v) => {
646 match self.get(v) {
648 Some(Ty::TyList(replacement)) => replacement.clone(),
649 Some(_) => l.clone(), None => l.clone(),
651 }
652 }
653 TyList::Append(xs, ys) => {
654 TyList::append(self.apply_ty_list(xs), self.apply_ty_list(ys))
655 }
656 }
657 }
658
659 #[must_use]
661 pub fn compose(&self, other: &Self) -> Self {
662 let mut result = Self::new();
663 for (k, v) in &other.mapping {
664 result.mapping.insert(*k, self.apply(v));
665 }
666 for (k, v) in &self.mapping {
667 result.mapping.entry(*k).or_insert_with(|| v.clone());
668 }
669 result
670 }
671
672 #[must_use]
674 pub fn is_empty(&self) -> bool {
675 self.mapping.is_empty()
676 }
677}
678
679#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
689pub enum PrimTy {
690 I32,
692 I64,
694 U32,
696 U64,
698 F32,
700 F64,
702 Char,
704 Addr,
706}
707
708impl PrimTy {
709 #[must_use]
711 pub const fn size_bytes(self) -> usize {
712 match self {
713 Self::I32 | Self::U32 | Self::F32 => 4,
714 Self::I64 | Self::U64 | Self::F64 | Self::Addr => 8,
715 Self::Char => 4, }
717 }
718
719 #[must_use]
721 pub const fn alignment(self) -> usize {
722 self.size_bytes()
723 }
724
725 #[must_use]
727 pub const fn name(self) -> &'static str {
728 match self {
729 Self::I32 => "Int32#",
730 Self::I64 => "Int#",
731 Self::U32 => "Word32#",
732 Self::U64 => "Word#",
733 Self::F32 => "Float#",
734 Self::F64 => "Double#",
735 Self::Char => "Char#",
736 Self::Addr => "Addr#",
737 }
738 }
739
740 #[must_use]
742 pub const fn is_signed_int(self) -> bool {
743 matches!(self, Self::I32 | Self::I64)
744 }
745
746 #[must_use]
748 pub const fn is_unsigned_int(self) -> bool {
749 matches!(self, Self::U32 | Self::U64)
750 }
751
752 #[must_use]
754 pub const fn is_float(self) -> bool {
755 matches!(self, Self::F32 | Self::F64)
756 }
757
758 #[must_use]
760 pub const fn is_numeric(self) -> bool {
761 matches!(
762 self,
763 Self::I32 | Self::I64 | Self::U32 | Self::U64 | Self::F32 | Self::F64
764 )
765 }
766}
767
768impl std::fmt::Display for PrimTy {
769 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
770 write!(f, "{}", self.name())
771 }
772}
773
774#[derive(Clone, Debug, thiserror::Error, Serialize, Deserialize)]
776pub enum TypeError {
777 #[error("type mismatch: expected {expected}, found {found}")]
779 Mismatch {
780 expected: String,
782 found: String,
784 span: Span,
786 },
787
788 #[error("infinite type: {var} occurs in {ty}")]
790 OccursCheck {
791 var: String,
793 ty: String,
795 span: Span,
797 },
798
799 #[error("unbound type variable: {name}")]
801 UnboundVar {
802 name: String,
804 span: Span,
806 },
807
808 #[error("kind mismatch: expected {expected}, found {found}")]
810 KindMismatch {
811 expected: String,
813 found: String,
815 span: Span,
817 },
818
819 #[error("ambiguous type variable: {var}")]
821 Ambiguous {
822 var: String,
824 span: Span,
826 },
827}
828
829#[cfg(test)]
830mod tests {
831 use super::*;
832
833 #[test]
834 fn test_ty_free_vars() {
835 let a = TyVar::new_star(0);
836 let b = TyVar::new_star(1);
837
838 let ty = Ty::fun(Ty::Var(a.clone()), Ty::Var(b.clone()));
840 let vars = ty.free_vars();
841 assert_eq!(vars.len(), 2);
842 assert!(vars.contains(&a));
843 assert!(vars.contains(&b));
844 }
845
846 #[test]
847 fn test_subst_apply() {
848 let a = TyVar::new_star(0);
849 let int_con = TyCon::new(unsafe { Symbol::from_raw(0) }, Kind::Star);
851
852 let mut subst = Subst::new();
853 subst.insert(&a, Ty::Con(int_con.clone()));
854
855 let ty = Ty::fun(Ty::Var(a.clone()), Ty::Var(a));
856 let result = subst.apply(&ty);
857
858 match result {
859 Ty::Fun(from, to) => {
860 assert!(matches!(*from, Ty::Con(_)));
861 assert!(matches!(*to, Ty::Con(_)));
862 }
863 _ => panic!("expected function type"),
864 }
865 }
866
867 #[test]
868 fn test_scheme_mono() {
869 let int_con = TyCon::new(unsafe { Symbol::from_raw(0) }, Kind::Star);
871 let scheme = Scheme::mono(Ty::Con(int_con));
872 assert!(scheme.is_mono());
873 }
874
875 #[test]
876 fn test_prim_ty_properties() {
877 assert_eq!(PrimTy::I64.size_bytes(), 8);
878 assert_eq!(PrimTy::I32.size_bytes(), 4);
879 assert_eq!(PrimTy::F64.size_bytes(), 8);
880 assert_eq!(PrimTy::F32.size_bytes(), 4);
881
882 assert!(PrimTy::I64.is_signed_int());
883 assert!(PrimTy::I32.is_signed_int());
884 assert!(!PrimTy::U64.is_signed_int());
885
886 assert!(PrimTy::F64.is_float());
887 assert!(PrimTy::F32.is_float());
888 assert!(!PrimTy::I64.is_float());
889
890 assert!(PrimTy::I64.is_numeric());
891 assert!(PrimTy::F64.is_numeric());
892 assert!(!PrimTy::Char.is_numeric());
893 assert!(!PrimTy::Addr.is_numeric());
894 }
895
896 #[test]
897 fn test_prim_ty_names() {
898 assert_eq!(PrimTy::I64.name(), "Int#");
899 assert_eq!(PrimTy::F64.name(), "Double#");
900 assert_eq!(PrimTy::F32.name(), "Float#");
901 assert_eq!(PrimTy::Char.name(), "Char#");
902 }
903
904 #[test]
905 fn test_ty_prim_constructors() {
906 let int = Ty::int_prim();
907 assert!(int.is_prim());
908 assert_eq!(int.as_prim(), Some(PrimTy::I64));
909
910 let double = Ty::double_prim();
911 assert!(double.is_prim());
912 assert_eq!(double.as_prim(), Some(PrimTy::F64));
913
914 let unit = Ty::unit();
916 assert!(!unit.is_prim());
917 assert_eq!(unit.as_prim(), None);
918 }
919
920 #[test]
921 fn test_prim_ty_no_free_vars() {
922 let prim = Ty::int_prim();
924 assert!(prim.free_vars().is_empty());
925 }
926
927 #[test]
928 fn test_subst_prim_unchanged() {
929 let a = TyVar::new_star(0);
931 let mut subst = Subst::new();
932 subst.insert(&a, Ty::unit());
933
934 let prim = Ty::int_prim();
935 let result = subst.apply(&prim);
936 assert_eq!(result, prim);
937 }
938}