bhc_types/
nat.rs

1//! Type-level natural numbers.
2//!
3//! This module implements type-level natural numbers for shape-indexed tensors,
4//! enabling compile-time dimension checking per H26-SPEC Section 7.
5//!
6//! ## Overview
7//!
8//! Type-level naturals are used to express tensor shapes at the type level:
9//!
10//! ```text
11//! matmul :: Tensor '[m, k] Float -> Tensor '[k, n] Float -> Tensor '[m, n] Float
12//! ```
13//!
14//! The `k` dimensions must match, which is enforced at compile time.
15//!
16//! ## Representation
17//!
18//! - `Lit(n)` - A concrete natural number
19//! - `Var(v)` - A polymorphic dimension variable
20//! - `Add(a, b)` - Sum of two naturals (m + k)
21//! - `Mul(a, b)` - Product of two naturals (m * k)
22
23use serde::{Deserialize, Serialize};
24
25use crate::TyVar;
26
27/// A type-level natural number.
28///
29/// Represents dimensions in shape-indexed tensor types. Supports both
30/// concrete values (for static shapes) and variables (for polymorphic
31/// operations like `matmul`).
32#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
33pub enum TyNat {
34    /// A concrete natural number literal.
35    ///
36    /// Example: `1024` in `Tensor '[1024, 768] Float`
37    Lit(u64),
38
39    /// A type-level natural variable.
40    ///
41    /// Example: `n` in `Tensor '[n] Float`
42    Var(TyVar),
43
44    /// Addition of two naturals.
45    ///
46    /// Example: `m + k` for computing output dimensions
47    Add(Box<TyNat>, Box<TyNat>),
48
49    /// Multiplication of two naturals.
50    ///
51    /// Example: `m * k` for computing buffer sizes
52    Mul(Box<TyNat>, Box<TyNat>),
53}
54
55impl TyNat {
56    /// Creates a literal type-level natural.
57    #[must_use]
58    pub fn lit(n: u64) -> Self {
59        Self::Lit(n)
60    }
61
62    /// Creates a zero type-level natural.
63    #[must_use]
64    pub fn zero() -> Self {
65        Self::Lit(0)
66    }
67
68    /// Creates a one type-level natural.
69    #[must_use]
70    pub fn one() -> Self {
71        Self::Lit(1)
72    }
73
74    /// Creates an addition of two naturals.
75    #[must_use]
76    pub fn add(a: TyNat, b: TyNat) -> Self {
77        // Simplify if both are literals
78        match (&a, &b) {
79            (TyNat::Lit(x), TyNat::Lit(y)) => TyNat::Lit(x + y),
80            _ => Self::Add(Box::new(a), Box::new(b)),
81        }
82    }
83
84    /// Creates a multiplication of two naturals.
85    #[must_use]
86    pub fn mul(a: TyNat, b: TyNat) -> Self {
87        // Simplify if both are literals
88        match (&a, &b) {
89            (TyNat::Lit(x), TyNat::Lit(y)) => TyNat::Lit(x * y),
90            _ => Self::Mul(Box::new(a), Box::new(b)),
91        }
92    }
93
94    /// Returns true if this is a concrete literal.
95    #[must_use]
96    pub fn is_lit(&self) -> bool {
97        matches!(self, Self::Lit(_))
98    }
99
100    /// Returns the literal value if this is a concrete literal.
101    #[must_use]
102    pub fn as_lit(&self) -> Option<u64> {
103        match self {
104            Self::Lit(n) => Some(*n),
105            _ => None,
106        }
107    }
108
109    /// Returns true if this natural contains no variables.
110    #[must_use]
111    pub fn is_ground(&self) -> bool {
112        match self {
113            Self::Lit(_) => true,
114            Self::Var(_) => false,
115            Self::Add(a, b) | Self::Mul(a, b) => a.is_ground() && b.is_ground(),
116        }
117    }
118
119    /// Evaluates this natural if it is ground (contains no variables).
120    ///
121    /// Returns `None` if the natural contains variables.
122    #[must_use]
123    pub fn eval(&self) -> Option<u64> {
124        match self {
125            Self::Lit(n) => Some(*n),
126            Self::Var(_) => None,
127            Self::Add(a, b) => Some(a.eval()? + b.eval()?),
128            Self::Mul(a, b) => Some(a.eval()? * b.eval()?),
129        }
130    }
131
132    /// Collects all type variables occurring in this natural.
133    #[must_use]
134    pub fn free_vars(&self) -> Vec<TyVar> {
135        let mut vars = Vec::new();
136        self.collect_free_vars(&mut vars);
137        vars
138    }
139
140    fn collect_free_vars(&self, vars: &mut Vec<TyVar>) {
141        match self {
142            Self::Lit(_) => {}
143            Self::Var(v) => {
144                if !vars.contains(v) {
145                    vars.push(v.clone());
146                }
147            }
148            Self::Add(a, b) | Self::Mul(a, b) => {
149                a.collect_free_vars(vars);
150                b.collect_free_vars(vars);
151            }
152        }
153    }
154
155    /// Substitutes a type variable with a natural in this expression.
156    #[must_use]
157    pub fn subst(&self, var: &TyVar, replacement: &TyNat) -> Self {
158        match self {
159            Self::Lit(n) => Self::Lit(*n),
160            Self::Var(v) if v.id == var.id => replacement.clone(),
161            Self::Var(v) => Self::Var(v.clone()),
162            Self::Add(a, b) => Self::add(a.subst(var, replacement), b.subst(var, replacement)),
163            Self::Mul(a, b) => Self::mul(a.subst(var, replacement), b.subst(var, replacement)),
164        }
165    }
166}
167
168impl std::fmt::Display for TyNat {
169    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
170        match self {
171            Self::Lit(n) => write!(f, "{n}"),
172            Self::Var(v) => write!(f, "n{}", v.id),
173            Self::Add(a, b) => write!(f, "({a} + {b})"),
174            Self::Mul(a, b) => write!(f, "({a} * {b})"),
175        }
176    }
177}
178
179#[cfg(test)]
180mod tests {
181    use super::*;
182    use crate::Kind;
183
184    #[test]
185    fn test_literal_naturals() {
186        let n = TyNat::lit(42);
187        assert!(n.is_lit());
188        assert!(n.is_ground());
189        assert_eq!(n.as_lit(), Some(42));
190        assert_eq!(n.eval(), Some(42));
191    }
192
193    #[test]
194    fn test_variable_naturals() {
195        let v = TyVar::new(0, Kind::Nat);
196        let n = TyNat::Var(v);
197        assert!(!n.is_lit());
198        assert!(!n.is_ground());
199        assert_eq!(n.as_lit(), None);
200        assert_eq!(n.eval(), None);
201    }
202
203    #[test]
204    fn test_add_literals() {
205        let a = TyNat::lit(10);
206        let b = TyNat::lit(20);
207        let sum = TyNat::add(a, b);
208        // Should simplify to Lit(30)
209        assert_eq!(sum, TyNat::lit(30));
210    }
211
212    #[test]
213    fn test_add_with_variable() {
214        let v = TyVar::new(0, Kind::Nat);
215        let a = TyNat::Var(v.clone());
216        let b = TyNat::lit(10);
217        let sum = TyNat::add(a, b);
218        // Should not simplify
219        assert!(!sum.is_lit());
220        assert!(!sum.is_ground());
221    }
222
223    #[test]
224    fn test_mul_literals() {
225        let a = TyNat::lit(3);
226        let b = TyNat::lit(4);
227        let prod = TyNat::mul(a, b);
228        assert_eq!(prod, TyNat::lit(12));
229    }
230
231    #[test]
232    fn test_free_vars() {
233        let v1 = TyVar::new(0, Kind::Nat);
234        let v2 = TyVar::new(1, Kind::Nat);
235        let n = TyNat::add(TyNat::Var(v1.clone()), TyNat::Var(v2.clone()));
236        let vars = n.free_vars();
237        assert_eq!(vars.len(), 2);
238        assert!(vars.contains(&v1));
239        assert!(vars.contains(&v2));
240    }
241
242    #[test]
243    fn test_substitution() {
244        let v = TyVar::new(0, Kind::Nat);
245        let n = TyNat::add(TyNat::Var(v.clone()), TyNat::lit(5));
246        let result = n.subst(&v, &TyNat::lit(10));
247        assert_eq!(result.eval(), Some(15));
248    }
249
250    #[test]
251    fn test_display() {
252        assert_eq!(format!("{}", TyNat::lit(42)), "42");
253        let v = TyVar::new(0, Kind::Nat);
254        assert_eq!(format!("{}", TyNat::Var(v)), "n0");
255    }
256}
bhc_types/nat.rs

bhc_types/
nat.rs
