morok-ir 0.1.0-alpha.2

Intermediate representation for the Morok ML compiler
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56

//! Pattern matching infrastructure for UOp graphs.
//!
//! This module provides pattern matching using `SimplifiedPatternMatcher`, which
//! uses closures that do inline Rust pattern matching. The `patterns!` macro
//! generates closures with native `match` expressions for O(1) OpKey dispatch.

pub mod helpers;
pub mod simplified;

use crate::UOp;
use std::sync::Arc;

// =============================================================================
// RewriteResult - Result of pattern matching
// =============================================================================

/// Result of applying a pattern rewrite.
#[derive(Debug, Clone)]
pub enum RewriteResult {
    /// Pattern didn't match or rewrite function declined to rewrite
    NoMatch,
    /// Pattern matched and returned a replacement UOp
    Rewritten(Arc<UOp>),
    /// Pattern matched and indicates bottom-up gate (Tinygrad's BottomUpGate)
    /// This signals that children should be processed before proceeding
    Gate(Arc<UOp>),
}

// =============================================================================
// Pattern Exports
// =============================================================================

pub use helpers::{const_matches, is_any_const, is_neg_one, is_nonzero, is_one, is_zero, try_const};
pub use simplified::{PatternClosure, SimplifiedPatternMatcher};

/// Type alias for backwards compatibility.
pub type TypedPatternMatcher<C = ()> = SimplifiedPatternMatcher<C>;

// =============================================================================
// Matcher Trait - Unified interface for pattern matchers
// =============================================================================

/// Trait for pattern matchers used by the rewrite engine.
///
/// This trait provides a unified interface for pattern matching,
/// allowing the rewrite engine to work with different matcher implementations.
pub trait Matcher<C> {
    /// Attempt to rewrite a UOp using registered patterns.
    fn rewrite(&self, uop: &Arc<UOp>, ctx: &mut C) -> RewriteResult;
}

impl<C, M: Matcher<C> + ?Sized> Matcher<C> for &M {
    fn rewrite(&self, uop: &Arc<UOp>, ctx: &mut C) -> RewriteResult {
        (**self).rewrite(uop, ctx)
    }
}