bend/fun/transform/
resolve_refs.rs

1use crate::{
2  diagnostics::Diagnostics,
3  fun::{Ctx, Name, Pattern, Term},
4  maybe_grow,
5};
6use std::collections::{HashMap, HashSet};
7
8#[derive(Debug, Clone)]
9pub struct ReferencedMainErr;
10
11impl Ctx<'_> {
12  /// Decides if names inside a term belong to a Var or to a Ref.
13  /// Converts `Term::Var(nam)` into `Term::Ref(nam)` when the name
14  /// refers to a function definition and there is no variable in
15  /// scope shadowing that definition.
16  ///
17  /// Precondition: Refs are encoded as vars, Constructors are resolved.
18  ///
19  /// Postcondition: Refs are encoded as refs, with the correct def id.
20  pub fn resolve_refs(&mut self) -> Result<(), Diagnostics> {
21    let def_names =
22      self.book.defs.keys().cloned().chain(self.book.hvm_defs.keys().cloned()).collect::<HashSet<_>>();
23    for (def_name, def) in &mut self.book.defs {
24      for rule in def.rules.iter_mut() {
25        let mut scope = HashMap::new();
26
27        for name in rule.pats.iter().flat_map(Pattern::binds) {
28          push_scope(name.as_ref(), &mut scope);
29        }
30
31        let res =
32          rule.body.resolve_refs(&def_names, self.book.entrypoint.as_ref(), &mut scope, &mut self.info);
33        self.info.take_rule_err(res, def_name.clone());
34      }
35    }
36
37    self.info.fatal(())
38  }
39}
40
41impl Term {
42  pub fn resolve_refs<'a>(
43    &'a mut self,
44    def_names: &HashSet<Name>,
45    main: Option<&Name>,
46    scope: &mut HashMap<&'a Name, usize>,
47    info: &mut Diagnostics,
48  ) -> Result<(), String> {
49    maybe_grow(move || {
50      match self {
51        Term::Var { nam } => {
52          if is_var_in_scope(nam, scope) {
53            // If the variable is actually a reference to main, don't swap and return an error.
54            if let Some(main) = main {
55              if nam == main {
56                return Err("Main definition can't be referenced inside the program.".to_string());
57              }
58            }
59
60            // If the variable is actually a reference to a function, swap the term.
61            if def_names.contains(nam) {
62              *self = Term::r#ref(nam);
63            }
64          }
65        }
66        Term::Def { def, nxt } => {
67          for rule in def.rules.iter_mut() {
68            let mut scope = HashMap::new();
69
70            for name in rule.pats.iter().flat_map(Pattern::binds) {
71              push_scope(name.as_ref(), &mut scope);
72            }
73
74            let res = rule.body.resolve_refs(def_names, main, &mut scope, info);
75            info.take_rule_err(res, def.name.clone());
76          }
77          nxt.resolve_refs(def_names, main, scope, info)?;
78        }
79        _ => {
80          for (child, binds) in self.children_mut_with_binds() {
81            for bind in binds.clone() {
82              push_scope(bind.as_ref(), scope);
83            }
84            child.resolve_refs(def_names, main, scope, info)?;
85            for bind in binds.rev() {
86              pop_scope(bind.as_ref(), scope);
87            }
88          }
89        }
90      }
91      Ok(())
92    })
93  }
94}
95
96fn push_scope<'a>(name: Option<&'a Name>, scope: &mut HashMap<&'a Name, usize>) {
97  if let Some(name) = name {
98    let var_scope = scope.entry(name).or_default();
99    *var_scope += 1;
100  }
101}
102
103fn pop_scope<'a>(name: Option<&'a Name>, scope: &mut HashMap<&'a Name, usize>) {
104  if let Some(name) = name {
105    let var_scope = scope.entry(name).or_default();
106    *var_scope -= 1;
107  }
108}
109
110fn is_var_in_scope<'a>(name: &'a Name, scope: &HashMap<&'a Name, usize>) -> bool {
111  match scope.get(name) {
112    Some(entry) => *entry == 0,
113    None => true,
114  }
115}