Witch of Git - mukan/blob - src/lib.rs

   1 use dyn_clone::{self, DynClone};
   2 use im::HashMap;
   3 use std::fmt;
   4
   5 pub mod term;
   6 pub use term::{Term, Var};
   7
   8 #[derive(Debug, Clone, Default)]
   9 pub struct State {
  10     fresh: u32,
  11     subst: Subst,
  12 }
  13
  14 pub fn delay<G: Goal>(thunk: impl Clone + Fn() -> G + 'static) -> BoxedGoal {
  15     BoxedGoal::new(move |state: State| thunk().search(state))
  16 }
  17
  18 #[derive(Clone)]
  19 pub struct BoxedGoal {
  20     inner: Box<dyn Goal<Stream = Box<dyn Iterator<Item = State>>>>,
  21 }
  22
  23 impl BoxedGoal {
  24     fn new(goal: impl Goal + Clone + 'static) -> Self {
  25         BoxedGoal {
  26             inner: Box::new(BoxGoalInner(goal)),
  27         }
  28     }
  29 }
  30
  31 impl Goal for BoxedGoal {
  32     type Stream = Box<dyn Iterator<Item = State>>;
  33     fn search(&self, state: State) -> Self::Stream {
  34         self.inner.search(state)
  35     }
  36 }
  37
  38 #[derive(Clone)]
  39 struct BoxGoalInner<G>(G);
  40
  41 impl<G> Goal for BoxGoalInner<G>
  42 where
  43     G: Goal + Clone,
  44     G::Stream: 'static,
  45 {
  46     type Stream = Box<dyn Iterator<Item = State>>;
  47     fn search(&self, state: State) -> Self::Stream {
  48         Box::new(self.0.search(state))
  49     }
  50 }
  51
  52 pub trait Goal: DynClone {
  53     type Stream: Iterator<Item = State>;
  54     fn search(&self, state: State) -> Self::Stream;
  55 }
  56
  57 dyn_clone::clone_trait_object!(Goal<Stream = Box<dyn Iterator<Item = State>>>);
  58
  59 pub fn call_fresh<G: Goal>(block: impl Clone + Fn(Var) -> G) -> impl Goal + Clone {
  60     move |mut state: State| {
  61         let var = Var(state.fresh);
  62         state.fresh += 1;
  63         block(var).search(state)
  64     }
  65 }
  66
  67 #[macro_export]
  68 macro_rules! fresh {
  69     (() => $e:expr) => { $e };
  70     (($v:ident $(, $vs:ident)* $(,)?) => $e:expr) => {
  71         call_fresh(move |$v| $crate::fresh!(($($vs),*) => $e))
  72     };
  73 }
  74
  75 #[macro_export]
  76 macro_rules! all {
  77     () => { |state: State| std::iter::once(state) };
  78     ($e:expr) => { $e };
  79     ($e:expr $(, $es:expr)* $(,)?) => {
  80         and($e, all!($($es),*))
  81     }
  82 }
  83
  84 #[macro_export]
  85 macro_rules! any {
  86     () => { |state: State| std::iter::empty() };
  87     ($e:expr) => { $e };
  88     ($e:expr $(, $es:expr)* $(,)?) => {
  89         or($e, any!($($es),*))
  90     }
  91 }
  92
  93 pub fn eq(u: impl Into<Term>, v: impl Into<Term>) -> impl Goal + Clone {
  94     let u = u.into();
  95     let v = v.into();
  96     move |state: State| state.unify(&u, &v).into_iter()
  97 }
  98
  99 pub fn and(u: impl Goal + Clone, v: impl Goal + Clone) -> impl Goal + Clone {
 100     move |state: State| Bind {
 101         a: u.search(state),
 102         b: v.clone(),
 103         current: None,
 104         done: false,
 105     }
 106 }
 107
 108 pub fn or(u: impl Goal + Clone, v: impl Goal + Clone) -> impl Goal + Clone {
 109     move |state: State| Plus {
 110         a: u.search(state.clone()),
 111         b: v.search(state),
 112         which: Which::A,
 113     }
 114 }
 115
 116 impl State {
 117     pub fn new() -> Self {
 118         Self::default()
 119     }
 120
 121     pub fn with_vars(n: u32) -> (Vec<Var>, Self) {
 122         let vars = (0..n).map(|v| Var(v)).collect();
 123         (
 124             vars,
 125             Self {
 126                 fresh: n + 1,
 127                 ..Default::default()
 128             },
 129         )
 130     }
 131
 132     pub fn unify(&self, u: &Term, v: &Term) -> Option<State> {
 133         self.subst.unify(u, v).map(|subst| State {
 134             fresh: self.fresh,
 135             subst,
 136         })
 137     }
 138
 139     pub fn eval_var(&self, var: Var) -> Term {
 140         let var = Term::Var(var);
 141         self.eval(&var)
 142     }
 143
 144     pub fn eval(&self, term: &Term) -> Term {
 145         match self.subst.walk(&term) {
 146             Term::Pair(u, v) => Term::from((self.eval(&*u), self.eval(&*v))),
 147             term => term.clone(),
 148         }
 149     }
 150 }
 151
 152 #[derive(Debug, Clone, Default)]
 153 struct Subst(HashMap<Var, Term>);
 154
 155 impl Subst {
 156     fn walk<'a>(&'a self, mut term: &'a Term) -> &'a Term {
 157         while let Term::Var(var) = term {
 158             match self.0.get(&var) {
 159                 Some(t) => term = t,
 160                 None => break,
 161             }
 162         }
 163         term
 164     }
 165
 166     fn extended(&self, var: Var, term: Term) -> Self {
 167         Subst(self.0.update(var, term))
 168     }
 169
 170     fn unify(&self, u: &Term, v: &Term) -> Option<Subst> {
 171         let u = self.walk(&u);
 172         let v = self.walk(&v);
 173         match (u, v) {
 174             (u, v) if u == v => Some(self.clone()),
 175             (Term::Var(u), term) => Some(self.clone().extended(*u, term.clone())),
 176             (term, Term::Var(v)) => Some(self.clone().extended(*v, term.clone())),
 177             (Term::Pair(u0, u1), Term::Pair(v0, v1)) => self
 178                 .unify(&*u0, &*v0)
 179                 .and_then(|s| s.unify(&*u1, &*v1)),
 180             _ => None,
 181         }
 182     }
 183 }
 184
 185 pub fn reify(state: State) -> HashMap<Var, Term> {
 186     let mut result = HashMap::new();
 187     for &v in state.subst.0.keys() {
 188         let var = Term::Var(v);
 189         let term = state.subst.walk(&var);
 190         result.insert(v, state.eval(term));
 191     }
 192     result
 193 }
 194
 195 pub fn reify_vars(vars: &[Var], state: State) -> HashMap<Var, Term> {
 196     let mut result = HashMap::new();
 197     for &v in vars {
 198         let var = Term::Var(v);
 199         let term = state.subst.walk(&var);
 200         result.insert(v, state.eval(term));
 201     }
 202     result
 203 }
 204
 205 impl<F, S> Goal for F
 206 where
 207     F: Clone + Fn(State) -> S,
 208     S: Iterator<Item = State>,
 209 {
 210     type Stream = S;
 211     fn search(&self, state: State) -> S {
 212         self(state)
 213     }
 214 }
 215
 216 enum Which {
 217     A,
 218     B,
 219     OnlyA,
 220     OnlyB,
 221     Done,
 222 }
 223
 224 struct Plus<S, R> {
 225     a: S,
 226     b: R,
 227     which: Which,
 228 }
 229
 230 impl<S, R, T> Iterator for Plus<S, R>
 231 where
 232     S: Iterator<Item = T>,
 233     R: Iterator<Item = T>,
 234     T: fmt::Display,
 235 {
 236     type Item = T;
 237     fn next(&mut self) -> Option<T> {
 238         match self.which {
 239             Which::A => match self.a.next() {
 240                 None => {
 241                     self.which = Which::OnlyB;
 242                     self.next()
 243                 }
 244                 item => {
 245                     self.which = Which::B;
 246                     item
 247                 }
 248             },
 249             Which::B => match self.b.next() {
 250                 None => {
 251                     self.which = Which::OnlyA;
 252                     self.next()
 253                 }
 254                 item => {
 255                     self.which = Which::A;
 256                     item
 257                 }
 258             },
 259             Which::OnlyA => match self.a.next() {
 260                 None => {
 261                     self.which = Which::Done;
 262                     None
 263                 }
 264                 item => item,
 265             },
 266             Which::OnlyB => match self.b.next() {
 267                 None => {
 268                     self.which = Which::Done;
 269                     None
 270                 }
 271                 item => item,
 272             },
 273             Which::Done => None,
 274         }
 275     }
 276 }
 277
 278 struct Bind<S: Iterator<Item = State>, R: Goal> {
 279     a: S,
 280     b: R,
 281     current: Option<R::Stream>,
 282     done: bool,
 283 }
 284
 285 impl<S, R> Iterator for Bind<S, R>
 286 where
 287     S: Iterator<Item = State>,
 288     R: Goal,
 289 {
 290     type Item = State;
 291
 292     fn next(&mut self) -> Option<Self::Item> {
 293         if self.done {
 294             return None;
 295         }
 296         match &mut self.current {
 297             None => {
 298                 if let Some(state) = self.a.next() {
 299                     self.current = Some(self.b.search(state));
 300                     self.next()
 301                 } else {
 302                     self.done = true;
 303                     None
 304                 }
 305             }
 306             Some(stream) => match stream.next() {
 307                 None => {
 308                     self.current = None;
 309                     self.next()
 310                 }
 311                 item => item,
 312             },
 313         }
 314     }
 315 }
 316
 317 impl fmt::Display for Subst {
 318     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 319         fmt.debug_map()
 320             .entries(self.0.iter().map(|x| (&x.0, &x.1)))
 321             .finish()
 322     }
 323 }
 324
 325 impl fmt::Display for State {
 326     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 327         write!(fmt, "#{} {}", self.fresh, self.subst)
 328     }
 329 }