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

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