1 files changed, 63 insertions, 1 deletions

diff --git a/generic/regc_nfa.c b/generic/regc_nfa.c
index dd26e64..48f56a9 100644
--- a/generic/regc_nfa.c
+++ b/generic/regc_nfa.c
@@ -61,11 +61,12 @@ struct nfa *parent;		/* NULL if primary NFA */
 	nfa->nstates = 0;
 	nfa->cm = cm;
 	nfa->v = v;
+	nfa->size = 0;
 	nfa->bos[0] = nfa->bos[1] = COLORLESS;
 	nfa->eos[0] = nfa->eos[1] = COLORLESS;
+	nfa->parent = parent;
 	nfa->post = newfstate(nfa, '@');	/* number 0 */
 	nfa->pre = newfstate(nfa, '>');		/* number 1 */
-	nfa->parent = parent;
 
 	nfa->init = newstate(nfa);		/* may become invalid later */
 	nfa->final = newstate(nfa);
@@ -88,6 +89,57 @@ struct nfa *parent;		/* NULL if primary NFA */
 }
 
 /*
+ - too_many_states - checks if the max states exceeds the compile-time value
+ ^ static int too_many_states(struct nfa *);
+ */
+static int
+too_many_states(nfa)
+struct nfa *nfa;
+{
+	struct nfa *parent = nfa->parent;
+	size_t sz = nfa->size;
+	while (parent != NULL) {
+		sz = parent->size;
+		parent = parent->parent;
+	}
+	if (sz > REG_MAX_STATES)
+		return 1;
+	return 0;
+}
+
+/*
+ - increment_size - increases the tracked size of the NFA and its parents.
+ ^ static void increment_size(struct nfa *);
+ */
+static void
+increment_size(nfa)
+struct nfa *nfa;
+{
+	struct nfa *parent = nfa->parent;
+	nfa->size++;
+	while (parent != NULL) {
+		parent->size++;
+		parent = parent->parent;
+	}
+}
+
+/*
+ - decrement_size - increases the tracked size of the NFA and its parents.
+ ^ static void decrement_size(struct nfa *);
+ */
+static void
+decrement_size(nfa)
+struct nfa *nfa;
+{
+	struct nfa *parent = nfa->parent;
+	nfa->size--;
+	while (parent != NULL) {
+		parent->size--;
+		parent = parent->parent;
+	}
+}
+
+/*
  - freenfa - free an entire NFA
  ^ static VOID freenfa(struct nfa *);
  */
@@ -123,6 +175,11 @@ struct nfa *nfa;
 {
 	struct state *s;
 
+	if (too_many_states(nfa)) {
+		/* XXX: add specific error for this */
+		NERR(REG_ETOOBIG);
+		return NULL;
+	}
 	if (nfa->free != NULL) {
 		s = nfa->free;
 		nfa->free = s->next;
@@ -154,6 +211,8 @@ struct nfa *nfa;
 	}
 	s->prev = nfa->slast;
 	nfa->slast = s;
+	/* Track the current size and the parent size */
+	increment_size(nfa);
 	return s;
 }
 
@@ -221,6 +280,7 @@ struct state *s;
 	s->prev = NULL;
 	s->next = nfa->free;	/* don't delete it, put it on the free list */
 	nfa->free = s;
+	decrement_size(nfa);
 }
 
 /*
@@ -651,6 +711,8 @@ struct state *stmp;		/* s's duplicate, or NULL */
 
 	for (a = s->outs; a != NULL && !NISERR(); a = a->outchain) {
 		duptraverse(nfa, a->to, (struct state *)NULL);
+		if (NISERR())
+			break;
 		assert(a->to->tmp != NULL);
 		cparc(nfa, a, s->tmp, a->to->tmp);
 	}