From 31492d0e2fcae3d43d6b1ae28fa4e6048dee3bef Mon Sep 17 00:00:00 2001
From: Vladimir Panteleev <vladimir@thecybershadow.net>
Date: Thu, 18 Mar 2010 21:17:29 +0200
Subject: [PATCH] Adding AbcTraitsVisitor and AsTraitsVisitor

---
 abcfile.d   | 30 ++++++++++++++++++++++++++++++
 asprogram.d | 48 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 78 insertions(+)

diff --git a/abcfile.d b/abcfile.d
index a3d8d3f..5005255 100644
--- a/abcfile.d
+++ b/abcfile.d
@@ -1788,3 +1788,33 @@ final:
 		writeU30(v.varName);
 	}
 }
+
+class ABCTraitsVisitor
+{
+	ABCFile abc;
+
+	this(ABCFile abc)
+	{
+		this.abc = abc;
+	}
+
+	final void run()
+	{
+		foreach (ref v; abc.instances)
+			visitTraits(v.traits);
+		foreach (ref v; abc.classes)
+			visitTraits(v.traits);
+		foreach (ref v; abc.scripts)
+			visitTraits(v.traits);
+		foreach (ref v; abc.bodies)
+			visitTraits(v.traits);
+	}
+
+	final void visitTraits(ABCFile.TraitsInfo[] traits)
+	{
+		foreach (ref trait; traits)
+			visitTrait(trait);
+	}
+
+	abstract void visitTrait(ref ABCFile.TraitsInfo trait);
+}
diff --git a/asprogram.d b/asprogram.d
index 466efd7..ba180be 100644
--- a/asprogram.d
+++ b/asprogram.d
@@ -1255,3 +1255,51 @@ final:
 		return r;
 	}
 }
+
+class ASTraitsVisitor
+{
+	ASProgram as;
+
+	this(ASProgram as)
+	{
+		this.as = as;
+	}
+
+	void run()
+	{
+		foreach (ref v; as.scripts)
+			visitTraits(v.traits);
+	}
+
+	final void visitTraits(ASProgram.Trait[] traits)
+	{
+		foreach (ref trait; traits)
+			visitTrait(trait);
+	}
+
+	void visitTrait(ref ASProgram.Trait trait)
+	{
+		switch (trait.kind)
+		{
+			case TraitKind.Slot:
+			case TraitKind.Const:
+				break;
+			case TraitKind.Class:
+				visitTraits(trait.vClass.vclass.traits);
+				visitTraits(trait.vClass.vclass.instance.traits);
+				break;
+			case TraitKind.Function:
+				if (trait.vFunction.vfunction.vbody)
+					visitTraits(trait.vFunction.vfunction.vbody.traits);
+				break;
+			case TraitKind.Method:
+			case TraitKind.Getter:
+			case TraitKind.Setter:
+				if (trait.vMethod.vmethod.vbody)
+					visitTraits(trait.vMethod.vmethod.vbody.traits);
+				break;
+			default:
+				throw new Exception("Unknown trait kind");
+		}
+	}
+}