zhc_ir 0.1.7

Graph-based intermediate representation framework with dialect support
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
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117

use std::hash::Hash;

use zhc_utils::{small::SmallVec, svec};

use crate::{
    Dialect, DialectInstructionSet, Format, FormatContext, IR, OpId, OpRef, Signature, ValId, sig,
    visualization::layoutlang::{LayoutDialect, LayoutTypeSystem},
};

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct OpContent {
    pub args: SmallVec<String>,
    pub returns: SmallVec<String>,
    pub call: String,
    pub comment: Option<String>,
}

impl OpContent {
    pub fn from_op<'ir, D: Dialect>(opref: &OpRef<'ir, D>, ctx: &FormatContext) -> Self {
        OpContent {
            args: opref
                .get_args_iter()
                .map(|a| a.fmt_to_string(ctx))
                .collect(),
            returns: opref
                .get_returns_iter()
                .map(|a| a.fmt_to_string(ctx))
                .collect(),
            call: opref.fmt_to_string(ctx),
            comment: opref.comment.clone(),
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum LayoutInstructionSet {
    Operation {
        opid: OpId,
        op: OpContent,
        args: SmallVec<ValId>,
        returns: SmallVec<ValId>,
    },
    Dummy {
        valid: ValId,
    },
    Group {
        ir: IR<LayoutDialect>,
        name: String,
    },
    GroupInput {
        pos: u16,
        valid: ValId,
    },
    GroupOutput {
        pos: u16,
        valid: ValId,
    },
}

impl Format for LayoutInstructionSet {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>, ctx: &crate::FormatContext) -> std::fmt::Result {
        match self {
            LayoutInstructionSet::Operation { opid, .. } => write!(f, "operation<{opid}>"),
            LayoutInstructionSet::Dummy { valid, .. } => write!(f, "dummy<{valid}>"),
            LayoutInstructionSet::Group { ir, name, .. } => {
                let inner_ctx = ctx.with_prefix("    ").with_next_nested_prefix();
                writeln!(f, "group<\"{}\"> {{", name)?;
                Format::fmt(ir, f, &inner_ctx)?;
                write!(f, "\n{}}}", ctx.prefix())
            }
            LayoutInstructionSet::GroupInput { pos, .. } => write!(f, "group_input<{pos}>"),
            LayoutInstructionSet::GroupOutput { pos, .. } => write!(f, "group_output<{pos}>"),
        }
    }
}

impl DialectInstructionSet for LayoutInstructionSet {
    type TypeSystem = LayoutTypeSystem;

    fn get_signature(&self) -> crate::Signature<Self::TypeSystem> {
        match self {
            LayoutInstructionSet::Operation { op, .. } => Signature(
                svec![LayoutTypeSystem::Value; op.args.len()],
                svec![LayoutTypeSystem::Value; op.returns.len()],
            ),
            LayoutInstructionSet::Dummy { .. } => {
                sig![(LayoutTypeSystem::Value) -> (LayoutTypeSystem::Value)]
            }
            LayoutInstructionSet::Group { ir, .. } => {
                let n_inputs = ir
                    .walk_ops_linear()
                    .filter(|op| {
                        matches!(
                            op.get_instruction(),
                            LayoutInstructionSet::GroupInput { .. }
                        )
                    })
                    .count();
                let n_outputs = ir
                    .walk_ops_linear()
                    .filter(|op| {
                        matches!(
                            op.get_instruction(),
                            LayoutInstructionSet::GroupOutput { .. }
                        )
                    })
                    .count();
                Signature(
                    svec![LayoutTypeSystem::Value; n_inputs],
                    svec![LayoutTypeSystem::Value; n_outputs],
                )
            }
            LayoutInstructionSet::GroupInput { .. } => sig![() -> (LayoutTypeSystem::Value)],
            LayoutInstructionSet::GroupOutput { .. } => sig![(LayoutTypeSystem::Value) -> ()],
        }
    }
}