From 29323e4f741d71f7ae585d8533dc1e086ad40654 Mon Sep 17 00:00:00 2001
From: Thomas Krijnen <t.krijnen@gmail.com>
Date: Tue, 1 Aug 2017 19:07:46 +0200
Subject: [PATCH] Fix recursion in traverse()

---
 src/ifcparse/IfcParse.cpp | 24 +++++++++++++++++++++++-
 test/tests.py             |  2 ++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/src/ifcparse/IfcParse.cpp b/src/ifcparse/IfcParse.cpp
index 71bba0b599..f1558b85a8 100644
--- a/src/ifcparse/IfcParse.cpp
+++ b/src/ifcparse/IfcParse.cpp
@@ -1361,6 +1361,24 @@ bool IfcFile::Init(IfcParse::IfcSpfStream* s) {
 	return true;
 }
 
+class traversal_visitor {
+private:
+	std::set<IfcUtil::IfcBaseClass*>& visited_;
+	IfcEntityList::ptr& list_;
+	int level_;
+	int max_level_;
+
+public:
+	traversal_visitor(std::set<IfcUtil::IfcBaseClass*>& visited, IfcEntityList::ptr& list, int level, int max_level)
+		: visited_(visited)
+		, list_(list)
+		, level_(level)
+		, max_level_(max_level)
+	{}
+
+	void operator()(IfcUtil::IfcBaseClass* inst);
+};
+
 void traverse_(IfcUtil::IfcBaseClass* instance, std::set<IfcUtil::IfcBaseClass*>& visited, IfcEntityList::ptr list, int level, int max_level) {
 	if (visited.find(instance) != visited.end()) {
 		return;
@@ -1370,10 +1388,14 @@ void traverse_(IfcUtil::IfcBaseClass* instance, std::set<IfcUtil::IfcBaseClass*>
 
 	if (level >= max_level && max_level > 0) return;
 
-	add_to_instance_list_visitor visit(list);
+	traversal_visitor visit(visited, list, level + 1, max_level);
 	apply_individual_instance_visitor(instance->entity).apply(visit);
 }
 
+void traversal_visitor::operator()(IfcUtil::IfcBaseClass* inst) {
+	traverse_(inst, visited_, list_, level_, max_level_);
+}
+
 IfcEntityList::ptr IfcParse::traverse(IfcUtil::IfcBaseClass* instance, int max_level) {
 	std::set<IfcUtil::IfcBaseClass*> visited;
 	IfcEntityList::ptr return_value(new IfcEntityList);
diff --git a/test/tests.py b/test/tests.py
index 358df9f7b0..df0d69d9ac 100644
--- a/test/tests.py
+++ b/test/tests.py
@@ -37,6 +37,8 @@ assert f.by_guid("28pa2ppDf1IA$BaQrvAf48").is_a("IfcProject")
 assert f.createIfcCartesianPoint((0., 0., 0.)).is_a("IfcCartesianPoint")
 assert f.by_type("IfcProject")[0].is_a("IfcProject")
 assert f.traverse(f[16])[-1].is_a("IFCSIUNIT")
+assert len(f.traverse(f[35], 1)) == 2
+assert len(f.traverse(f[35])) == 3
 assert f[16] in f.get_inverse(f[15])
 assert f[16].UnitComponent is not None
 f.remove(f[15])