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
118
119
120

use zhc_utils::{graphics::Frame, small::SmallVec};

use crate::{AnnIR, OpMap, ValMap, visualization::LayoutDialect};

use super::*;

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CompositionVariable {
    Op {
        sol: VariableCell,
        args: SmallVec<VariableCell>,
        rets: SmallVec<VariableCell>,
        body: VariableCell,
        comment: Option<VariableCell>,
    },
    Dummy {
        sol: VariableCell,
    },
    GroupInput {
        sol: VariableCell,
    },
    GroupOutput {
        sol: VariableCell,
    },
    Group {
        sol: VariableCell,
        inputs: SmallVec<VariableCell>,
        outputs: SmallVec<VariableCell>,
        maps: (OpMap<CompositionVariable>, ValMap<()>),
    },
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CompositionSolution {
    Op {
        sol: Solution,
        args: SmallVec<Solution>,
        rets: SmallVec<Solution>,
        body: Solution,
        comment: Option<Solution>,
    },
    Dummy {
        sol: Solution,
    },
    GroupInput {
        sol: Solution,
    },
    GroupOutput {
        sol: Solution,
    },
    Group {
        sol: Solution,
        inputs: SmallVec<Solution>,
        outputs: SmallVec<Solution>,
        maps: (OpMap<CompositionSolution>, ValMap<()>),
    },
}

impl CompositionSolution {
    /// Returns the overall frame for this element.
    pub fn get_frame(&self) -> Frame {
        match self {
            CompositionSolution::Op { sol, .. } => sol.frame.clone(),
            CompositionSolution::Dummy { sol } => sol.frame.clone(),
            CompositionSolution::GroupInput { sol } => sol.frame.clone(),
            CompositionSolution::GroupOutput { sol } => sol.frame.clone(),
            CompositionSolution::Group { sol, .. } => sol.frame.clone(),
        }
    }
}

fn resolve(input: OpMap<CompositionVariable>) -> OpMap<CompositionSolution> {
    input.map(|v| match v {
        CompositionVariable::Op {
            sol,
            args,
            rets,
            body,
            comment,
        } => CompositionSolution::Op {
            sol: sol.get_solution(),
            args: args.into_iter().map(|a| a.get_solution()).collect(),
            rets: rets.into_iter().map(|a| a.get_solution()).collect(),
            body: body.get_solution(),
            comment: comment.map(|a| a.get_solution()),
        },
        CompositionVariable::Dummy { sol } => CompositionSolution::Dummy {
            sol: sol.get_solution(),
        },
        CompositionVariable::GroupInput { sol } => CompositionSolution::GroupInput {
            sol: sol.get_solution(),
        },
        CompositionVariable::GroupOutput { sol } => CompositionSolution::GroupOutput {
            sol: sol.get_solution(),
        },
        CompositionVariable::Group {
            sol,
            inputs,
            outputs,
            maps,
        } => CompositionSolution::Group {
            sol: sol.get_solution(),
            inputs: inputs.into_iter().map(|a| a.get_solution()).collect(),
            outputs: outputs.into_iter().map(|a| a.get_solution()).collect(),
            maps: (resolve(maps.0), maps.1),
        },
    })
}

pub fn turn_to_solution(
    ir: AnnIR<'_, LayoutDialect, CompositionVariable, ()>,
) -> AnnIR<'_, LayoutDialect, CompositionSolution, ()> {
    let AnnIR {
        ir,
        op_annotations,
        val_annotations,
    } = ir;
    let op_annotations = resolve(op_annotations);
    AnnIR::new(ir, op_annotations, val_annotations)
}