Remove recursion when colorizing Transaction

MetacoSA · Oct 3, 2014 · a8d0490 · a8d0490
1 parent d0cab1a
commit a8d0490
Show file tree

Hide file tree

Showing 4 changed files with 140 additions and 12 deletions.
diff --git a/NBitcoin/CachedTransactionRepository.cs b/NBitcoin/CachedTransactionRepository.cs
@@ -0,0 +1,43 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+
+namespace NBitcoin
+{
+	public class CachedTransactionRepository : ITransactionRepository
+	{
+		ITransactionRepository _Inner;
+		Dictionary<uint256, Transaction> _Transactions = new Dictionary<uint256, Transaction>();
+		public CachedTransactionRepository(ITransactionRepository inner)
+		{
+			if(inner == null)
+				throw new ArgumentNullException("inner");
+			_Inner = inner;
+		}
+		#region ITransactionRepository Members
+
+		public Transaction Get(uint256 txId)
+		{
+			Transaction result = null;
+			if(!_Transactions.TryGetValue(txId, out result))
+			{
+				result = _Inner.Get(txId);
+				_Transactions.Add(txId, result);
+			}
+			return result;
+		}
+
+		public void Put(uint256 txId, Transaction tx)
+		{
+			if(!_Transactions.ContainsKey(txId))
+				_Transactions.Add(txId, tx);
+			else
+				_Transactions[txId] = tx;
+			_Inner.Put(txId, tx);
+		}
+
+		#endregion
+	}
+}
diff --git a/NBitcoin/NBitcoin.csproj b/NBitcoin/NBitcoin.csproj
@@ -72,6 +72,7 @@
     </Compile>
     <Compile Include="CachedBlockProvider.cs" />
     <Compile Include="CachedNoSqlRepository.cs" />
+    <Compile Include="CachedTransactionRepository.cs" />
     <Compile Include="CoinsView.cs" />
     <Compile Include="Crypto\DeterministicECDSA.cs" />
     <Compile Include="DataDirectory.cs" />
@@ -82,6 +83,7 @@
     <Compile Include="NoSqlTransactionRepository.cs" />
     <Compile Include="ObjectStream.cs" />
     <Compile Include="OpenAsset\Asset.cs" />
+    <Compile Include="OpenAsset\CachedColoredTransactionRepository.cs" />
     <Compile Include="OpenAsset\ColoredTransaction.cs" />
     <Compile Include="OpenAsset\Extensions.cs" />
     <Compile Include="OpenAsset\IColoredTransactionRepository.cs" />

diff --git a/NBitcoin/OpenAsset/CachedColoredTransactionRepository.cs b/NBitcoin/OpenAsset/CachedColoredTransactionRepository.cs
@@ -0,0 +1,53 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+
+namespace NBitcoin.OpenAsset
+{
+	public class CachedColoredTransactionRepository : IColoredTransactionRepository
+	{
+		IColoredTransactionRepository _Inner;
+		ITransactionRepository _InnerTransactionRepository;
+		Dictionary<uint256, ColoredTransaction> _ColoredTransactions = new Dictionary<uint256, ColoredTransaction>();
+		public CachedColoredTransactionRepository(IColoredTransactionRepository inner)
+		{
+			if(inner == null)
+				throw new ArgumentNullException("inner");
+			_Inner = inner;
+			_InnerTransactionRepository = new CachedTransactionRepository(inner.Transactions);
+		}
+		#region IColoredTransactionRepository Members
+
+		public ITransactionRepository Transactions
+		{
+			get
+			{
+				return _InnerTransactionRepository;
+			}
+		}
+
+		public ColoredTransaction Get(uint256 txId)
+		{
+			ColoredTransaction result = null;
+			if(!_ColoredTransactions.TryGetValue(txId, out result))
+			{
+				result = _Inner.Get(txId);
+				_ColoredTransactions.Add(txId, result);
+			}
+			return result;
+		}
+
+		public void Put(uint256 txId, ColoredTransaction tx)
+		{
+			if(!_ColoredTransactions.ContainsKey(txId))
+				_ColoredTransactions.Add(txId, tx);
+			else
+				_ColoredTransactions[txId] = tx;
+			_Inner.Put(txId, tx);
+		}
+
+		#endregion
+	}
+}
diff --git a/NBitcoin/OpenAsset/ColoredTransaction.cs b/NBitcoin/OpenAsset/ColoredTransaction.cs
@@ -59,7 +59,7 @@ public override string ToString()
 			if(Asset == null)
 				return "[" + Index + "]";
 			else
-				return "[" + Index + "] " + Asset; 
+				return "[" + Index + "] " + Asset;
 		}
 	}
 	public class ColoredTransaction : IBitcoinSerializable
@@ -84,7 +84,47 @@ public static ColoredTransaction FetchColors(uint256 txId, Transaction tx, IColo
 			var result = repo.Get(txId);
 			if(result != null)
 				return result;
+			repo = new CachedColoredTransactionRepository(repo);
+			HashSet<uint256> invalidColored = new HashSet<uint256>();
+			Stack<Tuple<uint256, Transaction>> ancestors = new Stack<Tuple<uint256, Transaction>>();
+			ancestors.Push(Tuple.Create(txId, tx));
+			while(ancestors.Count != 0)
+			{
+				var peek = ancestors.Peek();
+				txId = peek.Item1;
+				tx = peek.Item2;
+				bool isComplete = true;
+				if(!tx.HasWellFormedColoredMarker() && ancestors.Count != 1)
+				{
+					invalidColored.Add(txId);
+					ancestors.Pop();
+					continue;
+				}
+
+				for(int i = 0 ; i < tx.Inputs.Count ; i++)
+				{
+					var txin = tx.Inputs[i];
+					if(repo.Get(txin.PrevOut.Hash) == null && !invalidColored.Contains(txin.PrevOut.Hash))
+					{
+						var prevTx = repo.Transactions.Get(txin.PrevOut.Hash);
+						if(prevTx == null)
+							throw new TransactionNotFoundException("Transaction " + txin.PrevOut.Hash + " not found in transaction repository", txId);
+						ancestors.Push(Tuple.Create(txin.PrevOut.Hash, prevTx));
+						isComplete = false;
+					}
+				}
+				if(isComplete)
+				{
+					FetchColorsWithAncestorsSolved(txId, tx, repo);
+					ancestors.Pop();
+				}
+			}
+
+			return FetchColorsWithAncestorsSolved(txId, tx, repo);
+		}
 
+		private static ColoredTransaction FetchColorsWithAncestorsSolved(uint256 txId, Transaction tx, IColoredTransactionRepository repo)
+		{
 			ColoredTransaction colored = new ColoredTransaction();
 
 			Queue<ColoredEntry> previousAssetQueue = new Queue<ColoredEntry>();
@@ -93,17 +133,7 @@ public static ColoredTransaction FetchColors(uint256 txId, Transaction tx, IColo
 				var txin = tx.Inputs[i];
 				var prevColored = repo.Get(txin.PrevOut.Hash);
 				if(prevColored == null)
-				{
-					var prevTx = repo.Transactions.Get(txin.PrevOut.Hash);
-					if(prevTx == null)
-						throw new TransactionNotFoundException("Transaction " + txin.PrevOut.Hash + " not found in transaction repository", txId);
-					if(!prevTx.HasWellFormedColoredMarker())
-					{
-						continue;
-					}
-					prevColored = FetchColors(txin.PrevOut.Hash, prevTx, repo);
-				}
-
+					continue;
 				var prevAsset = prevColored.GetColoredEntry(txin.PrevOut.N);
 				if(prevAsset != null)
 				{