Merge pull request #10 from nerai/dictbuilder

improve DictBuilder

Author: dscheg

ZstdNet/DictBuilder.cs

@@ -1,24 +1,31 @@
 using System;
 using System.Collections.Generic;
+using System.IO;
 using System.Linq;
 using size_t = System.UIntPtr;
 
 namespace ZstdNet
 {
 	public static class DictBuilder
 	{
-		public static byte[] TrainFromBuffer(ICollection<byte[]> samples, int dictCapacity = DefaultDictCapacity)
+		public static byte[] TrainFromBuffer(IEnumerable<byte[]> samples, int dictCapacity = DefaultDictCapacity)
 		{
-			var samplesBuffer = samples.SelectMany(sample => sample).ToArray();
-			var samplesSizes = samples.Select(sample => (size_t)sample.Length).ToArray();
-			var dictBuffer = new byte[dictCapacity];
-			var dictSize = (int)ExternMethods.ZDICT_trainFromBuffer(dictBuffer, (size_t)dictCapacity, samplesBuffer, samplesSizes, (uint)samples.Count).EnsureZdictSuccess();
+			var ms = new MemoryStream ();
+			var samplesSizes = new List<size_t> ();
+			foreach (var sample in samples) {
+				samplesSizes.Add ((size_t) sample.Length);
+				ms.Write (sample, 0, sample.Length);
+			}
+			var samplesBuffer = ms.ToArray ();
 
-			if (dictCapacity == dictSize)
-				return dictBuffer;
-			var result = new byte[dictSize];
-			Array.Copy(dictBuffer, result, dictSize);
-			return result;
+			var dictBuffer = new byte[dictCapacity];
+			var dictSize = (int)ExternMethods
+				.ZDICT_trainFromBuffer(dictBuffer, (size_t)dictCapacity, samplesBuffer, samplesSizes.ToArray(), (uint)samplesSizes.Count)
+				.EnsureZdictSuccess();
+			if (dictCapacity != dictSize) {
+				Array.Resize<byte> (ref dictBuffer, dictSize);
+			}
+			return dictBuffer;
 		}
 
 		public const int DefaultDictCapacity = 112640; // Used by zstd utility by default