1 files changed, 43 insertions, 3 deletions

diff --git a/src/generator.hpp b/src/generator.hpp
index 47ef6bf..bed048f 100644
--- a/src/generator.hpp
+++ b/src/generator.hpp
@@ -44,10 +44,11 @@ public:
 	};
 
 private:
-	std::vector<std::pair<item_t, std::string>> rules;
+	std::vector<std::tuple<item_t, std::vector<std::string>, std::string>> rules;
 	std::map<item_t, size_t> rule_ids;
 	std::multimap<std::string, size_t> nonterms;
 
+	std::set<std::string> nonterminals;
 	std::set<symbol_t> terminals;
 
 	std::multimap<symbol_t, item_t> items;
@@ -59,6 +60,9 @@ private:
 
 	std::set<size_t> shift_conflicts;
 
+	std::map<std::string, std::string> nonterm_types;
+	std::map<symbol_t, std::pair<std::string, std::string>> term_types;
+
 	void close_set(std::set<item_t> *set);
 	std::set<item_t> get_set(const std::string &nonterm);
 
@@ -90,6 +94,40 @@ private:
 	void generate_itemsets();
 
 public:
+	const std::string & get_nonterm_type(const std::string &sym) const {
+		static const std::string empty;
+
+		auto it = nonterm_types.find(sym);
+		if (it == nonterm_types.end())
+			return empty;
+		else
+			return it->second;
+	}
+
+	const std::pair<std::string, std::string> & get_term_type(const symbol_t &sym) const {
+		static const std::pair<std::string, std::string> empty;
+
+		auto it = term_types.find(sym);
+		if (it == term_types.end())
+			return empty;
+		else
+			return it->second;
+	}
+
+	const std::string & get_type(const symbol_t &sym) const {
+		switch (sym.get_type()) {
+		case SYMBOL_TYPE_NONTERM:
+			return get_nonterm_type(sym.get_value());
+
+		default:
+			return get_term_type(sym).first;
+		}
+	}
+
+	const std::set<std::string> & get_nonterminals() const {
+		return nonterminals;
+	}
+
 	const std::set<symbol_t> & get_terminals() const {
 		return terminals;
 	}
@@ -98,7 +136,7 @@ public:
 		return itemsets.size();
 	}
 
-	const std::vector<std::pair<item_t, std::string>> & get_rules() const {
+	const std::vector<std::tuple<item_t, std::vector<std::string>, std::string>> & get_rules() const {
 		return rules;
 	}
 
@@ -114,7 +152,9 @@ public:
 		return gotos;
 	}
 
-	generator_t(const std::vector<std::pair<item_t, std::string>> &rules0);
+	generator_t(const std::vector<std::tuple<item_t, std::vector<std::string>, std::string>> &rules0,
+		    const std::map<std::string, std::string> &nonterm_types0,
+		    const std::map<symbol_t, std::pair<std::string, std::string>> &term_types0);
 };
 
 }