<commit_before>#include "type.hpp"

#include <vector>
#include <sstream>

#include "ast.hpp"
#include "tool.hpp"

#include "maybe.hpp"
#include "package.hpp"

namespace type {

static const bool debug = false;
  
static const auto make_constant = [](const char* name, kind::any k=kind::term()) {
  return make_ref<constant>(constant{name, k});
};

  
static const auto make_variable = [](std::size_t level, kind::any k=kind::term()) {
  return make_ref<variable>(variable{level, k});
};

template<class Func>
static void iter_rows(mono self, Func func);
  
  
  // constants
  const mono unit = make_constant("unit");
  const mono boolean = make_constant("boolean");
  const mono integer = make_constant("integer");
  const mono real = make_constant("real");
  

  const mono func =
    make_constant("->", kind::term() >>= kind::term() >>= kind::term());

  const mono io =
    make_constant("io", kind::term() >>= kind::term());
  
  // records
  const mono record =
    make_constant("record", kind::row() >>= kind::term());

  const mono empty = make_constant("{}", kind::row());
  
  mono ext(symbol attr) {
    static const kind::any k = kind::term() >>= kind::row() >>= kind::row();  
    static std::map<symbol, ref<constant>> table;
    const std::string name = attr.get() + std::string(":");
  
    auto it = table.emplace(attr, make_constant(name.c_str(), k));
    return it.first->second;
  }

  const mono ty = make_constant("type", kind::term() >>= kind::term());
  
  
state& state::def(symbol name, mono t) {
  if(!vars->locals.emplace(name, generalize(t)).second) {
    throw error("redefined variable " + tool::quote(name.get()));
  }
  
  return *this;
}
  


mono mono::operator()(mono arg) const {
  return make_ref<application>(*this, arg);
}

mono operator>>=(mono from, mono to) {
  return mono(func)(from)(to);
}



application::application(mono ctor, mono arg)
  : ctor(ctor),
    arg(arg) {
  const kind::any k = ctor.kind();
  if(auto c = k.get<kind::constructor>()) {
    if(*c->from != arg.kind()) {
      std::clog << poly{{}, ctor} << " :: " << ctor.kind()
                << " " << poly{{}, arg} << " :: " << arg.kind()
                << std::endl;
      throw kind::error("argument does not have the expected kind");
    }
  } else {
    throw kind::error("type constructor must have constructor kind");
  }
}


struct ostream_visitor {
  using type = void;
  
  using map_type = std::map<ref<variable>, std::size_t>;
  map_type& map;

  void operator()(const ref<constant>& self, std::ostream& out,
                  const poly::forall_type& forall, bool parens) const {
    out << self->name;
  }

  void operator()(const ref<variable>& self, std::ostream& out,
                  const poly::forall_type& forall, bool parens) const {
    auto it = map.emplace(self, map.size()).first;
    if(forall.find(self) != forall.end()) {
      out << "'";
    } else {
      out << "!";
    }
    
    out << char('a' + it->second);

    static const bool extra = false;
    if(debug && extra) {
      out << "("
          << std::hex << (long(self.get()) & 0xffff)
          << "::" << self->kind
          << ")";
    }
  }

  void operator()(const ref<application>& self, std::ostream& out,
                  const poly::forall_type& forall, bool parens) const {
    const kind::any k = mono(self).kind();
    
    if(self->ctor == record) {
      // skip record constructor altogether since row types have {} around
      self->arg.visit(*this, out, forall, true);
      return;
    }

    if(parens) {
      if(k == kind::row()) out << "{";
      else out << "(";
    }
    
    if(self->ctor == func) {
      // print function types in infix style
      self->arg.visit(*this, out, forall, true);
      out << " ";
      self->ctor.visit(*this, out, forall, false);
    } else if(mono(self).kind() == kind::row()) {
      // print row types: ctor is (attr: type) so we want parens
      self->ctor.visit(*this, out, forall, true);

      // skip trailing empty record type
      if(self->arg != empty) {
        out << " ";
        self->arg.visit(*this, out, forall, false);
      }
    } else {
      // regular types
      self->ctor.visit(*this, out, forall, false);
      out << " ";

      const auto func_rhs = [&] {
        if(auto a = self->ctor.get<app>()) {
          if((*a)->ctor == func) return true;
        }
        return false;
      };

      // don't put parens on function rhs
      self->arg.visit(*this, out, forall, !func_rhs());
    }

    if(parens) {
      if(k == kind::row()) out << "}";
      else out << ")";
    }
  }
  
};


struct show {
  ostream_visitor::map_type& map;
  const poly& p;

  show(ostream_visitor::map_type& map, const poly& p)
    : map(map), p(p) { }
  
  friend std::ostream& operator<<(std::ostream& out, const show& self) {
    self.p.type.visit(ostream_visitor{self.map}, out, self.p.forall, false);
    return out;
  }
};


  

std::ostream& operator<<(std::ostream& out, const poly& self) {
  ostream_visitor::map_type map;
  self.type.visit(ostream_visitor{map}, out, self.forall, false);
  return out;
}


  
struct kind_visitor {
  using type = kind::any;
  
  kind::any operator()(const ref<constant>& self) const {
    return self->kind;
  }

  kind::any operator()(const ref<variable>& self) const {
    return self->kind;
  }

  kind::any operator()(const ref<application>& self) const {
    const kind::any k = self->ctor.kind();
    if(auto c = k.get<kind::constructor>()) {
      return *c->to;
    } else {
      throw std::logic_error("derp");
    }
  }
      
};

kind::any mono::kind() const {
  return visit(kind_visitor());
}



  
  struct let {
    const ::list<ast::bind> defs;
    const ast::expr body;

    static symbol fix() { return "__fix__"; }
    
    // rewrite let as non-recursive let + fix **when defining functions**
    static let rewrite(const ast::let& self) {
      using ::list;
      const list<ast::bind> defs = map(self.defs, [](ast::bind self) {
          return self.value.match<ast::bind>
          ([&](const ast::abs& abs) { 
            return ast::bind{self.name,
                             ast::app{ast::var{fix()},
                                      ast::abs{self.name >>= list<ast::abs::arg>(),
                                          self.value}
                                      >>= list<ast::expr>()}};
          }, 
            [&](const ast::expr& expr) { return self; });
        });

      return {defs, *self.body};
    }
  };

struct infer_visitor {
  using type = mono;
  
  template<class T>
  mono operator()(const T& self, const ref<state>&) const {
    throw std::logic_error("infer unimplemented: " + tool::type_name(typeid(T)));
  }


  // var
  mono operator()(const ast::var& self, const ref<state>& s) const {
    try {
      return s->instantiate(s->vars->find(self.name));
    } catch(std::out_of_range& e) {
      throw error("unbound variable " + tool::quote(self.name.get()));
    }
  }


  // reconstruct actual type from reified type: ... -> (type 'a) yields 'a
  // note: t must be substituted
  static mono reconstruct(ref<state> s, mono t) {
    // std::clog << "reconstructing: " << s->generalize(t) << std::endl;
    
    if(auto self = t.get<app>()) {
      // std::clog << "ctor: " << s->generalize((*self)->ctor) << std::endl;
      
      if((*self)->ctor == ty) {
        return (*self)->arg;
      }
    }
    
    auto from = s->fresh();
    auto to = s->fresh();
    
    s->unify(from >>= to, t);
    return reconstruct(s, s->substitute(to));
  }
  
  static cst constructor(mono t) {
    if(auto self = t.get<app>()) {
      return constructor((*self)->ctor);
    }

    if(auto self = t.get<cst>()) {
      return *self;
    }

    // TODO restrictive?
    throw error("constructor must be a constant");
  }
  
  
  // abs
  mono operator()(const ast::abs& self, const ref<state>& s) const {

    // function scope
    const auto sub = scope(s);

    // construct function type
    const mono result = s->fresh();
    
    const mono res = foldr(result, self.args, [&](ast::abs::arg arg, mono tail) {
      const mono sig = arg.match<mono>
        ([&](symbol self) {
          // untyped arguments: trivial signature
          const mono a = sub->fresh();
          return a >>= a;
        },
          [&](ast::abs::typed self) {
            // TODO reconstruct type
            const mono t = infer(s, self.type);
            const mono r = reconstruct(s, s->substitute(t));
            
            const cst c = constructor(r);
            // std::clog << "ctor: " << c->name << std::endl;
            try {
              auto sig = s->sigs->at(c);
              // std::clog << "sig: " << sig << std::endl;
              return sub->instantiate(sig);
              
            } catch(std::out_of_range&) {
              // TODO wtf
              throw error("unknown signature " + tool::quote(c->name.get()));
            }
          });
      
      const mono outer = s->fresh();
      const mono inner = sub->fresh();
      
      sub->unify(outer >>= inner, sig);
      sub->def(arg.name(), inner);

      // std::clog << "inner: " << sub->vars->find(arg.name()) << std::endl;
      
      return outer >>= tail;
    });

    
    
    // infer lambda body with augmented environment
    s->unify(result, infer(sub, *self.body));
    
    return res;
  }


  // app
  mono operator()(const ast::app& self, const ref<state>& s) const {
    const mono func = infer(s, *self.func);
    const mono result = s->fresh();

    const mono sig = foldr(result, self.args, [&](ast::expr e, mono t) {
        return infer(s, e) >>= t;
      });

    s->unify(sig, func);
    return result;
  }


  // non-recursive let
  mono operator()(const let& self, const ref<state>& s) const {
    auto sub = scope(s);

    for(ast::bind def : self.defs) {
      sub->vars->locals.emplace(def.name, s->generalize(infer(s, def.value)));
    }
    
    return infer(sub, self.body);
  }
  
  // recursive let
  mono operator()(const ast::let& self, const ref<state>& s) const {
    auto sub = scope(s);

    const mono a = sub->fresh();
    sub->def(let::fix(), (a >>= a) >>= a);
    
    return operator()(let::rewrite(self), sub);
  }


  // sel
  mono operator()(const ast::sel& self, const ref<state>& s) const {<commit_msg>comments<commit_after>#include "type.hpp"

#include <vector>
#include <sstream>

#include "ast.hpp"
#include "tool.hpp"

#include "maybe.hpp"
#include "package.hpp"

namespace type {

static const bool debug = false;
  
static const auto make_constant = [](const char* name, kind::any k=kind::term()) {
  return make_ref<constant>(constant{name, k});
};

  
static const auto make_variable = [](std::size_t level, kind::any k=kind::term()) {
  return make_ref<variable>(variable{level, k});
};

template<class Func>
static void iter_rows(mono self, Func func);
  
  
  // constants
  const mono unit = make_constant("unit");
  const mono boolean = make_constant("boolean");
  const mono integer = make_constant("integer");
  const mono real = make_constant("real");
  

  const mono func =
    make_constant("->", kind::term() >>= kind::term() >>= kind::term());

  const mono io =
    make_constant("io", kind::term() >>= kind::term());
  
  // records
  const mono record =
    make_constant("record", kind::row() >>= kind::term());

  const mono empty = make_constant("{}", kind::row());
  
  mono ext(symbol attr) {
    static const kind::any k = kind::term() >>= kind::row() >>= kind::row();  
    static std::map<symbol, ref<constant>> table;
    const std::string name = attr.get() + std::string(":");
  
    auto it = table.emplace(attr, make_constant(name.c_str(), k));
    return it.first->second;
  }

  const mono ty = make_constant("type", kind::term() >>= kind::term());
  
  
state& state::def(symbol name, mono t) {
  if(!vars->locals.emplace(name, generalize(t)).second) {
    throw error("redefined variable " + tool::quote(name.get()));
  }
  
  return *this;
}
  


mono mono::operator()(mono arg) const {
  return make_ref<application>(*this, arg);
}

mono operator>>=(mono from, mono to) {
  return mono(func)(from)(to);
}



application::application(mono ctor, mono arg)
  : ctor(ctor),
    arg(arg) {
  const kind::any k = ctor.kind();
  if(auto c = k.get<kind::constructor>()) {
    if(*c->from != arg.kind()) {
      std::clog << poly{{}, ctor} << " :: " << ctor.kind()
                << " " << poly{{}, arg} << " :: " << arg.kind()
                << std::endl;
      throw kind::error("argument does not have the expected kind");
    }
  } else {
    throw kind::error("type constructor must have constructor kind");
  }
}


struct ostream_visitor {
  using type = void;
  
  using map_type = std::map<ref<variable>, std::size_t>;
  map_type& map;

  void operator()(const ref<constant>& self, std::ostream& out,
                  const poly::forall_type& forall, bool parens) const {
    out << self->name;
  }

  void operator()(const ref<variable>& self, std::ostream& out,
                  const poly::forall_type& forall, bool parens) const {
    auto it = map.emplace(self, map.size()).first;
    if(forall.find(self) != forall.end()) {
      out << "'";
    } else {
      out << "!";
    }
    
    out << char('a' + it->second);

    static const bool extra = false;
    if(debug && extra) {
      out << "("
          << std::hex << (long(self.get()) & 0xffff)
          << "::" << self->kind
          << ")";
    }
  }

  void operator()(const ref<application>& self, std::ostream& out,
                  const poly::forall_type& forall, bool parens) const {
    const kind::any k = mono(self).kind();
    
    if(self->ctor == record) {
      // skip record constructor altogether since row types have {} around
      self->arg.visit(*this, out, forall, true);
      return;
    }

    if(parens) {
      if(k == kind::row()) out << "{";
      else out << "(";
    }
    
    if(self->ctor == func) {
      // print function types in infix style
      self->arg.visit(*this, out, forall, true);
      out << " ";
      self->ctor.visit(*this, out, forall, false);
    } else if(mono(self).kind() == kind::row()) {
      // print row types: ctor is (attr: type) so we want parens
      self->ctor.visit(*this, out, forall, true);

      // skip trailing empty record type
      if(self->arg != empty) {
        out << " ";
        self->arg.visit(*this, out, forall, false);
      }
    } else {
      // regular types
      self->ctor.visit(*this, out, forall, false);
      out << " ";

      const auto func_rhs = [&] {
        if(auto a = self->ctor.get<app>()) {
          if((*a)->ctor == func) return true;
        }
        return false;
      };

      // don't put parens on function rhs
      self->arg.visit(*this, out, forall, !func_rhs());
    }

    if(parens) {
      if(k == kind::row()) out << "}";
      else out << ")";
    }
  }
  
};


struct show {
  ostream_visitor::map_type& map;
  const poly& p;

  show(ostream_visitor::map_type& map, const poly& p)
    : map(map), p(p) { }
  
  friend std::ostream& operator<<(std::ostream& out, const show& self) {
    self.p.type.visit(ostream_visitor{self.map}, out, self.p.forall, false);
    return out;
  }
};


  

std::ostream& operator<<(std::ostream& out, const poly& self) {
  ostream_visitor::map_type map;
  self.type.visit(ostream_visitor{map}, out, self.forall, false);
  return out;
}


  
struct kind_visitor {
  using type = kind::any;
  
  kind::any operator()(const ref<constant>& self) const {
    return self->kind;
  }

  kind::any operator()(const ref<variable>& self) const {
    return self->kind;
  }

  kind::any operator()(const ref<application>& self) const {
    const kind::any k = self->ctor.kind();
    if(auto c = k.get<kind::constructor>()) {
      return *c->to;
    } else {
      throw std::logic_error("derp");
    }
  }
      
};

kind::any mono::kind() const {
  return visit(kind_visitor());
}



  
  struct let {
    const ::list<ast::bind> defs;
    const ast::expr body;

    static symbol fix() { return "__fix__"; }
    
    // rewrite let as non-recursive let + fix **when defining functions**
    static let rewrite(const ast::let& self) {
      using ::list;
      const list<ast::bind> defs = map(self.defs, [](ast::bind self) {
          return self.value.match<ast::bind>
          ([&](const ast::abs& abs) { 
            return ast::bind{self.name,
                             ast::app{ast::var{fix()},
                                      ast::abs{self.name >>= list<ast::abs::arg>(),
                                          self.value}
                                      >>= list<ast::expr>()}};
          }, 
            [&](const ast::expr& expr) { return self; });
        });

      return {defs, *self.body};
    }
  };


  // rewrite app as nested unary applications
  static ast::app rewrite(const ast::app& self) {
    const ast::expr& init = *self.func;
    return foldl(init, self.args, [](ast::expr func, ast::expr arg) {
      return ast::app(func, arg >>= list<ast::expr>());
      }).cast<ast::app>();
  }

  
struct infer_visitor {
  using type = mono;
  
  template<class T>
  mono operator()(const T& self, const ref<state>&) const {
    throw std::logic_error("infer unimplemented: " + tool::type_name(typeid(T)));
  }


  // var
  mono operator()(const ast::var& self, const ref<state>& s) const {
    try {
      return s->instantiate(s->vars->find(self.name));
    } catch(std::out_of_range& e) {
      throw error("unbound variable " + tool::quote(self.name.get()));
    }
  }


  // reconstruct actual type from reified type: ... -> (type 'a) yields 'a
  // note: t must be substituted
  static mono reconstruct(ref<state> s, mono t) {
    // std::clog << "reconstructing: " << s->generalize(t) << std::endl;
    
    if(auto self = t.get<app>()) {
      // std::clog << "ctor: " << s->generalize((*self)->ctor) << std::endl;
      
      if((*self)->ctor == ty) {
        return (*self)->arg;
      }
    }
    
    auto from = s->fresh();
    auto to = s->fresh();
    
    s->unify(from >>= to, t);
    return reconstruct(s, s->substitute(to));
  }
  
  static cst constructor(mono t) {
    if(auto self = t.get<app>()) {
      return constructor((*self)->ctor);
    }

    if(auto self = t.get<cst>()) {
      return *self;
    }

    // TODO restrictive?
    throw error("constructor must be a constant");
  }
  
  
  // abs
  mono operator()(const ast::abs& self, const ref<state>& s) const {

    // function scope
    const auto sub = scope(s);

    // construct function type
    const mono result = s->fresh();
    
    const mono res = foldr(result, self.args, [&](ast::abs::arg arg, mono tail) {
      const mono sig = arg.match<mono>
        ([&](symbol self) {
          // untyped arguments: trivial signature
          const mono a = sub->fresh();
          return a >>= a;
        },
          [&](ast::abs::typed self) {
            // obtain reified type from annoatation
            const mono t = infer(s, self.type);

            // extract actual type from reified
            const mono r = reconstruct(s, s->substitute(t));

            // obtain type constructor from type
            const cst c = constructor(r);

            try {
              // fetch associated signature, if any
              auto sig = s->sigs->at(c);

              // instantiate signature
              return sub->instantiate(sig);
              
            } catch(std::out_of_range&) {
              throw error("unknown signature " + tool::quote(c->name.get()));
            }
          });
      
      const mono outer = s->fresh();
      const mono inner = sub->fresh();
      
      sub->unify(outer >>= inner, sig);
      sub->def(arg.name(), inner);

      // std::clog << "inner: " << sub->vars->find(arg.name()) << std::endl;
      
      return outer >>= tail;
    });

    
    
    // infer lambda body with augmented environment
    s->unify(result, infer(sub, *self.body));
    
    return res;
  }


  // app
  mono operator()(const ast::app& self, const ref<state>& s) const {
    // normalize application as unary
    const ast::app rw = rewrite(self);
    assert(size(rw.args) == 1);
    
    const mono func = infer(s, *rw.func);

    const mono arg = infer(s, rw.args->head);
    const mono ret = s->fresh();
    
    s->unify(arg >>= ret, func);
    
    return ret;
  }


  // non-recursive let
  mono operator()(const let& self, const ref<state>& s) const {
    auto sub = scope(s);

    for(ast::bind def : self.defs) {
      sub->vars->locals.emplace(def.name, s->generalize(infer(s, def.value)));
    }
    
    return infer(sub, self.body);
  }
  
  // recursive let
  mono operator()(const ast::let& self, const ref<state>& s) const {
    auto sub = scope(s);

    const mono a = sub->fresh();
    sub->def(let::fix(), (a >>= a) >>= a);
    
    return operator()(let::rewrite(self), sub);
  }


  // sel
  mono operator()(const ast::sel& self, const ref<state>& s) const {
{'fraction_non_alphanumeric': 0.17669172932330826, 'fraction_numerical': 0.0001427619682116684, 'mean_word_length': 2.405996758508914, 'pattern_counts': {'":': 4, '<': 235, '<?xml version=': 0, '>': 267, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7598310', 'n_tokens_mistral': 7063, 'n_tokens_neox': 6795, 'n_words': 1691}