2 files changed, 129 insertions, 9 deletions

diff --git a/src/program.c b/src/program.c
index d1fbbf3d..01a1db31 100644
--- a/src/program.c
+++ b/src/program.c
@@ -233,7 +233,8 @@ void colm_run_program( Program *prg, int argc, const char **argv )
 	prg->argv = 0;
 }
 
-Tree *colm_run_func( struct colm_program *prg, int frameId, const char **params, int paramCount )
+Tree *colm_run_func( struct colm_program *prg, int frameId,
+		const char **params, int paramCount )
 {
 	/* Make the arguments available to the program. */
 	prg->argc = 0;
@@ -282,22 +283,26 @@ Tree *colm_run_func( struct colm_program *prg, int frameId, const char **params,
 	return prg->returnVal;
 };
 
-int colm_delete_program( Program *prg )
+static void colm_clear_heap( Program *prg, Tree **sp )
 {
-	Tree **sp = prg->stackRoot;
-	int exitStatus = prg->exitStatus;
-
-	treeDownref( prg, sp, prg->returnVal );
-	clearGlobal( prg, sp );
-
 	/* Clear the heap. */
 	Kid *a = prg->heap;
 	while ( a != 0 ) {
 		Kid *next = a->next;
-		treeDownref( prg, sp, a->tree );
+		objectDownref( prg, sp, a->tree );
 		kidFree( prg, a );
 		a = next;
 	}
+}
+
+int colm_delete_program( Program *prg )
+{
+	Tree **sp = prg->stackRoot;
+	int exitStatus = prg->exitStatus;
+
+	treeDownref( prg, sp, prg->returnVal );
+	clearGlobal( prg, sp );
+	colm_clear_heap( prg, sp );
 
 	treeDownref( prg, sp, prg->trueVal );
 	treeDownref( prg, sp, prg->falseVal );
diff --git a/src/tree.c b/src/tree.c
index 0b2cdb4b..dc9d1379 100644
--- a/src/tree.c
+++ b/src/tree.c
@@ -1170,6 +1170,121 @@ void treeDownref( Program *prg, Tree **sp, Tree *tree )
 	}
 }
 
+/* We can't make recursive calls here since the tree we are freeing may be
+ * very large. Need the VM stack. */
+void objectFreeRec( Program *prg, Tree **sp, Tree *tree )
+{
+	Tree **top = vm_ptop();
+	LangElInfo *lelInfo;
+	long genericId;
+
+free_tree:
+	lelInfo = prg->rtd->lelInfo;
+	genericId = lelInfo[tree->id].genericId;
+	if ( genericId > 0 ) {
+		GenericInfo *generic = &prg->rtd->genericInfo[genericId];
+		if ( generic->type == GEN_LIST ) {
+			List *list = (List*) tree;
+			ListEl *el = list->head;
+			while ( el != 0 ) {
+				ListEl *next = el->next;
+				vm_push( el->value );
+				listElFree( prg, el );
+				el = next;
+			}
+			mapElFree( prg, (MapEl*)list );
+		}
+		else if ( generic->type == GEN_MAP ) {
+			Map *map = (Map*)tree;
+			MapEl *el = map->head;
+			while ( el != 0 ) {
+				MapEl *next = el->next;
+				vm_push( el->key );
+				vm_push( el->tree );
+				mapElFree( prg, el );
+				el = next;
+			}
+			mapElFree( prg, (MapEl*)map );
+		}
+		else if ( generic->type == GEN_PARSER ) {
+			Parser *parser = (Parser*)tree;
+			clearPdaRun( prg, sp, parser->pdaRun );
+			free( parser->pdaRun );
+			treeDownref( prg, sp, (Tree*)parser->input );
+			mapElFree( prg, (MapEl*)parser );
+		}
+		else {
+			assert(false);
+		}
+	}
+	else {
+		switch ( tree->id ) {
+		case LEL_ID_STR: {
+			Str *str = (Str*) tree;
+			stringFree( prg, str->value );
+			treeFree( prg, tree );
+			break;
+		}
+		case LEL_ID_BOOL:
+		case LEL_ID_INT: {
+			treeFree( prg, tree );
+			break;
+		}
+		case LEL_ID_PTR: {
+			treeFree( prg, tree );
+			break;
+		}
+		case LEL_ID_STREAM: {
+			Stream *stream = (Stream*)tree;
+			clearSourceStream( prg, sp, stream->in );
+			if ( stream->in->file != 0 )
+				fclose( stream->in->file );
+			else if ( stream->in->fd >= 0 )
+				close( stream->in->fd );
+			free( stream->in );
+			streamFree( prg, stream );
+			break;
+		}
+		default: { 
+			if ( tree->id != LEL_ID_IGNORE )
+				stringFree( prg, tree->tokdata );
+
+			/* Attributes and grammar-based children. */
+			Kid *child = tree->child;
+			while ( child != 0 ) {
+				Kid *next = child->next;
+				vm_push( child->tree );
+				kidFree( prg, child );
+				child = next;
+			}
+
+			treeFree( prg, tree );
+			break;
+		}}
+	}
+
+	/* Any trees to downref? */
+	while ( sp != top ) {
+		tree = vm_pop();
+		if ( tree != 0 ) {
+			assert( tree->refs > 0 );
+			tree->refs -= 1;
+			if ( tree->refs == 0 )
+				goto free_tree;
+		}
+	}
+}
+
+void objectDownref( Program *prg, Tree **sp, Tree *tree )
+{
+	if ( tree != 0 ) {
+		assert( tree->refs > 0 );
+		tree->refs -= 1;
+		if ( tree->refs == 0 )
+			objectFreeRec( prg, sp, tree );
+	}
+}
+
 /* Find the first child of a tree. */
 Kid *treeChild( Program *prg, const Tree *tree )
 {