aws_smt_ir/fold/
compose.rs

1// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// SPDX-License-Identifier: Apache-2.0
3use super::{Fold, Folder};
4use crate::{IConst, ICoreOp, ILet, IMatch, IOp, IQuantifier, IVar, Logic, IUF};
5use std::{
6    convert::{identity, Infallible},
7    marker::PhantomData,
8};
9
10pub trait Compose: Sized {
11    fn compose<After, M1, M2>(self, after: After) -> Composed<Self, After, M1, M2> {
12        TryComposed::new(self, after, identity, identity)
13    }
14
15    fn try_compose<After, MapErr1, MapErr2, Err, M1, M2>(
16        self,
17        after: After,
18        f: MapErr1,
19        g: MapErr2,
20    ) -> TryComposed<Self, After, MapErr1, MapErr2, Err, M1, M2> {
21        TryComposed::new(self, after, f, g)
22    }
23}
24
25impl<T> Compose for T {}
26
27type Composed<Before, After, M1, M2> = TryComposed<
28    Before,
29    After,
30    fn(Infallible) -> Infallible,
31    fn(Infallible) -> Infallible,
32    Infallible,
33    M1,
34    M2,
35>;
36
37pub struct TryComposed<Before, After, MapErr1, MapErr2, Err, M1, M2> {
38    before: Before,
39    after: After,
40    map_err1: MapErr1,
41    map_err2: MapErr2,
42    err: PhantomData<Err>,
43    m1: PhantomData<M1>,
44    m2: PhantomData<M2>,
45}
46
47impl<Before, After, MapErr1, MapErr2, Err, M1, M2>
48    TryComposed<Before, After, MapErr1, MapErr2, Err, M1, M2>
49{
50    const fn new(before: Before, after: After, map_err1: MapErr1, map_err2: MapErr2) -> Self {
51        Self {
52            before,
53            after,
54            map_err1,
55            map_err2,
56            err: PhantomData,
57            m1: PhantomData,
58            m2: PhantomData,
59        }
60    }
61}
62
63impl<L: Logic, Before, After, MapErr1, MapErr2, Err, M1, M2> Folder<L>
64    for TryComposed<Before, After, MapErr1, MapErr2, Err, M1, M2>
65where
66    Before: Folder<L, M1>,
67    Before::Output: Fold<L, After::Output, Output = After::Output>,
68    After: Folder<L, M2>,
69    MapErr1: FnMut(Before::Error) -> Err,
70    MapErr2: FnMut(After::Error) -> Err,
71{
72    type Output = After::Output;
73    type Error = Err;
74
75    fn fold_const(&mut self, constant: IConst) -> Result<Self::Output, Self::Error> {
76        let t = constant
77            .fold_with(&mut self.before)
78            .map_err(&mut self.map_err1)?;
79        t.fold_with(&mut self.after).map_err(&mut self.map_err2)
80    }
81
82    fn fold_var(&mut self, var: IVar<L::Var>) -> Result<Self::Output, Self::Error> {
83        let t = var
84            .fold_with(&mut self.before)
85            .map_err(&mut self.map_err1)?;
86        t.fold_with(&mut self.after).map_err(&mut self.map_err2)
87    }
88
89    fn fold_core_op(&mut self, op: ICoreOp<L>) -> Result<Self::Output, Self::Error> {
90        let t = op.fold_with(&mut self.before).map_err(&mut self.map_err1)?;
91        t.fold_with(&mut self.after).map_err(&mut self.map_err2)
92    }
93
94    fn fold_theory_op(&mut self, op: IOp<L>) -> Result<Self::Output, Self::Error> {
95        let t = op.fold_with(&mut self.before).map_err(&mut self.map_err1)?;
96        t.fold_with(&mut self.after).map_err(&mut self.map_err2)
97    }
98
99    fn fold_uninterpreted_func(&mut self, func: IUF<L>) -> Result<Self::Output, Self::Error> {
100        let t = func
101            .fold_with(&mut self.before)
102            .map_err(&mut self.map_err1)?;
103        t.fold_with(&mut self.after).map_err(&mut self.map_err2)
104    }
105
106    fn fold_let(&mut self, l: ILet<L>) -> Result<Self::Output, Self::Error> {
107        let t = l.fold_with(&mut self.before).map_err(&mut self.map_err1)?;
108        t.fold_with(&mut self.after).map_err(&mut self.map_err2)
109    }
110
111    fn fold_match(&mut self, m: IMatch<L>) -> Result<Self::Output, Self::Error> {
112        let t = m.fold_with(&mut self.before).map_err(&mut self.map_err1)?;
113        t.fold_with(&mut self.after).map_err(&mut self.map_err2)
114    }
115
116    fn fold_quantifier(&mut self, q: IQuantifier<L>) -> Result<Self::Output, Self::Error> {
117        let t = q.fold_with(&mut self.before).map_err(&mut self.map_err1)?;
118        t.fold_with(&mut self.after).map_err(&mut self.map_err2)
119    }
120}