BsonTrie.cs 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351
  1. /* Copyright 2010-present MongoDB Inc.
  2. *
  3. * Licensed under the Apache License, Version 2.0 (the "License");
  4. * you may not use this file except in compliance with the License.
  5. * You may obtain a copy of the License at
  6. *
  7. * http://www.apache.org/licenses/LICENSE-2.0
  8. *
  9. * Unless required by applicable law or agreed to in writing, software
  10. * distributed under the License is distributed on an "AS IS" BASIS,
  11. * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. * See the License for the specific language governing permissions and
  13. * limitations under the License.
  14. */
  15. using System;
  16. namespace MongoDB.Bson.IO
  17. {
  18. /// <summary>
  19. /// Represents a mapping from a set of UTF8 encoded strings to a set of elementName/value pairs, implemented as a trie.
  20. /// </summary>
  21. /// <typeparam name="TValue">The type of the BsonTrie values.</typeparam>
  22. public class BsonTrie<TValue>
  23. {
  24. // private fields
  25. private readonly BsonTrieNode<TValue> _root;
  26. // constructors
  27. /// <summary>
  28. /// Initializes a new instance of the BsonTrie class.
  29. /// </summary>
  30. public BsonTrie()
  31. {
  32. _root = new BsonTrieNode<TValue>(0);
  33. }
  34. // public properties
  35. /// <summary>
  36. /// Gets the root node.
  37. /// </summary>
  38. public BsonTrieNode<TValue> Root
  39. {
  40. get
  41. {
  42. return _root;
  43. }
  44. }
  45. // public methods
  46. /// <summary>
  47. /// Adds the specified elementName (after encoding as a UTF8 byte sequence) and value to the trie.
  48. /// </summary>
  49. /// <param name="elementName">The element name to add.</param>
  50. /// <param name="value">The value to add. The value can be null for reference types.</param>
  51. public void Add(string elementName, TValue value)
  52. {
  53. using var rentedSegment = Utf8Encodings.Strict.GetBytesUsingThreadStaticBuffer(elementName);
  54. var node = _root;
  55. foreach (var keyByte in rentedSegment.Segment)
  56. {
  57. var child = node.GetChild(keyByte);
  58. if (child == null)
  59. {
  60. child = new BsonTrieNode<TValue>(keyByte);
  61. node.AddChild(child);
  62. }
  63. node = child;
  64. }
  65. node.SetValue(elementName, value);
  66. }
  67. /// <summary>
  68. /// Gets the node associated with the specified element name.
  69. /// </summary>
  70. /// <param name="utf8">The element name.</param>
  71. /// <param name="node">
  72. /// When this method returns, contains the node associated with the specified element name, if the key is found;
  73. /// otherwise, null. This parameter is passed uninitialized.
  74. /// </param>
  75. /// <returns>True if the node was found; otherwise, false.</returns>
  76. public bool TryGetNode(ArraySegment<byte> utf8, out BsonTrieNode<TValue> node)
  77. {
  78. node = _root;
  79. for (var i = 0; node != null && i < utf8.Count; i++)
  80. {
  81. var keyByte = utf8.Array[utf8.Offset + i];
  82. node = node.GetChild(keyByte);
  83. }
  84. return node != null;
  85. }
  86. /// <summary>
  87. /// Tries to get the node associated with a name read from a stream.
  88. /// </summary>
  89. /// <param name="stream">The stream.</param>
  90. /// <param name="node">The node.</param>
  91. /// <returns>
  92. /// True if the node was found.
  93. /// If the node was found the stream is advanced over the name, otherwise
  94. /// the stream is repositioned to the beginning of the name.
  95. /// </returns>
  96. public bool TryGetNode(BsonStream stream, out BsonTrieNode<TValue> node)
  97. {
  98. var position = stream.Position;
  99. var utf8 = stream.ReadCStringBytes();
  100. if (TryGetNode(utf8, out node))
  101. {
  102. return true;
  103. }
  104. stream.Position = position;
  105. return false;
  106. }
  107. /// <summary>
  108. /// Gets the value associated with the specified element name.
  109. /// </summary>
  110. /// <param name="utf8">The element name.</param>
  111. /// <param name="value">
  112. /// When this method returns, contains the value associated with the specified element name, if the key is found;
  113. /// otherwise, the default value for the type of the value parameter. This parameter is passed uninitialized.
  114. /// </param>
  115. /// <returns>True if the value was found; otherwise, false.</returns>
  116. public bool TryGetValue(ArraySegment<byte> utf8, out TValue value)
  117. {
  118. BsonTrieNode<TValue> node;
  119. if (TryGetNode(utf8, out node) && node.HasValue)
  120. {
  121. value = node.Value;
  122. return true;
  123. }
  124. value = default(TValue);
  125. return false;
  126. }
  127. /// <summary>
  128. /// Gets the value associated with the specified element name.
  129. /// </summary>
  130. /// <param name="elementName">The element name.</param>
  131. /// <param name="value">
  132. /// When this method returns, contains the value associated with the specified element name, if the key is found;
  133. /// otherwise, the default value for the type of the value parameter. This parameter is passed uninitialized.
  134. /// </param>
  135. /// <returns>True if the value was found; otherwise, false.</returns>
  136. public bool TryGetValue(string elementName, out TValue value)
  137. {
  138. using var utf8 = Utf8Encodings.Strict.GetBytesUsingThreadStaticBuffer(elementName);
  139. return TryGetValue(utf8.Segment, out value);
  140. }
  141. }
  142. /// <summary>
  143. /// Represents a node in a BsonTrie.
  144. /// </summary>
  145. /// <typeparam name="TValue">The type of the BsonTrie values.</typeparam>
  146. public sealed class BsonTrieNode<TValue>
  147. {
  148. // private fields
  149. private readonly byte _keyByte;
  150. private string _elementName;
  151. private TValue _value;
  152. private BsonTrieNode<TValue> _onlyChild; // used when there is only one child
  153. private BsonTrieNode<TValue>[] _children; // used when there are two or more children
  154. private byte[] _childrenIndexes; // maps key bytes into indexes into the _children array
  155. private byte _minChildKeyByte; // key byte value of first element in _childrenIndexes
  156. // constructors
  157. internal BsonTrieNode(byte keyByte)
  158. {
  159. _keyByte = keyByte;
  160. }
  161. /// <summary>
  162. /// Gets whether this node has a value.
  163. /// </summary>
  164. public bool HasValue
  165. {
  166. get
  167. {
  168. return _elementName != null;
  169. }
  170. }
  171. /// <summary>
  172. /// Gets the element name for this node.
  173. /// </summary>
  174. public string ElementName
  175. {
  176. get
  177. {
  178. if (_elementName == null)
  179. {
  180. throw new InvalidOperationException("BsonTrieNode doesn't have a value.");
  181. }
  182. return _elementName;
  183. }
  184. }
  185. /// <summary>
  186. /// Gets the value for this node.
  187. /// </summary>
  188. public TValue Value
  189. {
  190. get
  191. {
  192. if (_elementName == null)
  193. {
  194. throw new InvalidOperationException("BsonTrieNode doesn't have a value.");
  195. }
  196. return _value;
  197. }
  198. }
  199. // public methods
  200. /// <summary>
  201. /// Gets the child of this node for a given key byte.
  202. /// </summary>
  203. /// <param name="keyByte">The key byte.</param>
  204. /// <returns>The child node if it exists; otherwise, null.</returns>
  205. public BsonTrieNode<TValue> GetChild(byte keyByte)
  206. {
  207. if (_onlyChild != null)
  208. {
  209. // optimization for nodes that have only one child
  210. if (_onlyChild._keyByte == keyByte)
  211. {
  212. return _onlyChild;
  213. }
  214. }
  215. else if (_children != null)
  216. {
  217. var index = (uint)((int)keyByte - _minChildKeyByte);
  218. if (index < _childrenIndexes.Length)
  219. {
  220. index = _childrenIndexes[index];
  221. if (index < _children.Length)
  222. {
  223. return _children[index];
  224. }
  225. }
  226. }
  227. return null;
  228. }
  229. // internal methods
  230. internal void AddChild(BsonTrieNode<TValue> child)
  231. {
  232. if (GetChild(child._keyByte) != null)
  233. {
  234. throw new ArgumentException("BsonTrieNode already contains a child with the same keyByte.");
  235. }
  236. if (_children != null)
  237. {
  238. // add a new child to the existing _children
  239. var children = new BsonTrieNode<TValue>[_children.Length + 1];
  240. Array.Copy(_children, children, _children.Length);
  241. children[children.Length - 1] = child;
  242. var childrenIndexes = _childrenIndexes;
  243. var minChildKeyByte = _minChildKeyByte;
  244. var maxChildKeyByte = _minChildKeyByte + _childrenIndexes.Length - 1;
  245. // if new keyByte doesn't fall within existing min/max range expand the range
  246. if (child._keyByte < minChildKeyByte)
  247. {
  248. // grow the indexes on the min side
  249. minChildKeyByte = child._keyByte;
  250. childrenIndexes = new byte[maxChildKeyByte - minChildKeyByte + 1];
  251. var sizeDelta = childrenIndexes.Length - _childrenIndexes.Length;
  252. for (var i = 0; i < sizeDelta; i++)
  253. {
  254. childrenIndexes[i] = 255;
  255. }
  256. Array.Copy(_childrenIndexes, 0, childrenIndexes, sizeDelta, _childrenIndexes.Length);
  257. }
  258. else if (child._keyByte > maxChildKeyByte)
  259. {
  260. // grow the indexes on the max side
  261. maxChildKeyByte = child._keyByte;
  262. childrenIndexes = new byte[maxChildKeyByte - minChildKeyByte + 1];
  263. Array.Copy(_childrenIndexes, 0, childrenIndexes, 0, _childrenIndexes.Length);
  264. for (var i = _childrenIndexes.Length; i < childrenIndexes.Length; i++)
  265. {
  266. childrenIndexes[i] = 255;
  267. }
  268. }
  269. childrenIndexes[child._keyByte - minChildKeyByte] = (byte)(children.Length - 1);
  270. _children = children;
  271. _childrenIndexes = childrenIndexes;
  272. _minChildKeyByte = minChildKeyByte;
  273. }
  274. else if (_onlyChild != null)
  275. {
  276. // switch from having an _onlyChild to having two _children
  277. var children = new BsonTrieNode<TValue>[2];
  278. children[0] = _onlyChild;
  279. children[1] = child;
  280. var minChildKeyByte = _onlyChild._keyByte;
  281. var maxChildKeyByte = child._keyByte;
  282. if (minChildKeyByte > maxChildKeyByte)
  283. {
  284. minChildKeyByte = child._keyByte;
  285. maxChildKeyByte = _onlyChild._keyByte;
  286. }
  287. var childrenIndexes = new byte[maxChildKeyByte - minChildKeyByte + 1];
  288. for (var i = 0; i < childrenIndexes.Length; i++)
  289. {
  290. childrenIndexes[i] = 255;
  291. }
  292. childrenIndexes[_onlyChild._keyByte - minChildKeyByte] = 0;
  293. childrenIndexes[child._keyByte - minChildKeyByte] = 1;
  294. _onlyChild = null;
  295. _children = children;
  296. _childrenIndexes = childrenIndexes;
  297. _minChildKeyByte = minChildKeyByte;
  298. }
  299. else
  300. {
  301. _onlyChild = child;
  302. }
  303. }
  304. internal void SetValue(string elementName, TValue value)
  305. {
  306. if (elementName == null)
  307. {
  308. throw new ArgumentNullException("elementName");
  309. }
  310. if (_elementName != null)
  311. {
  312. throw new InvalidOperationException("BsonTrieNode already has a value.");
  313. }
  314. _elementName = elementName;
  315. _value = value;
  316. }
  317. }
  318. }