chomsky-extract 0.0.1

E-Graph extraction and selection logic for Chomsky optimization framework
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
57
58
59
60
61

use chomsky_cost::{Cost, CostModel};
use chomsky_extract::{IKunExtractor, IKunTree};
use chomsky_uir::{EGraph, IKun};

struct MockCostModel;
impl CostModel for MockCostModel {
    fn cost(&self, enode: &IKun) -> Cost {
        match enode {
            IKun::Constant(_) => Cost {
                latency: 1.0,
                ..Default::default()
            },
            IKun::Map(_, _) => Cost {
                latency: 10.0,
                ..Default::default()
            },
            _ => Cost::default(),
        }
    }
}

#[test]
fn test_extraction_basic() {
    let egraph = EGraph::<IKun, ()>::new();
    let c = egraph.add(IKun::Constant(42));

    let extractor = IKunExtractor::new(&egraph, MockCostModel);
    let tree = extractor.extract(c);

    if let IKunTree::Constant(val) = tree {
        assert_eq!(val, 42);
    } else {
        panic!("Expected Constant");
    }
}

#[test]
fn test_extraction_nested() {
    let egraph = EGraph::<IKun, ()>::new();
    let f = egraph.add(IKun::Symbol("f".to_string()));
    let x = egraph.add(IKun::Symbol("x".to_string()));
    let map = egraph.add(IKun::Map(f, x));

    let extractor = IKunExtractor::new(&egraph, MockCostModel);
    let tree = extractor.extract(map);

    if let IKunTree::Map(f_tree, x_tree) = tree {
        if let IKunTree::Symbol(s) = *f_tree {
            assert_eq!(s, "f");
        } else {
            panic!("Expected Symbol f");
        }
        if let IKunTree::Symbol(s) = *x_tree {
            assert_eq!(s, "x");
        } else {
            panic!("Expected Symbol x");
        }
    } else {
        panic!("Expected Map");
    }
}