Memory-efficient string collection

Question

My goal is to create a memory-efficient (immutable) collection of strings. The imagined use-case is checking for valid words in a Scrabble-like game.

Wikimedia Commons has a pretty good summary of the data structure (author: Chkno):

(Image license: Creative Commons Attribution-Share Alike 3.0 Unported.)

For an input of { "tap", "taps", "top", "tops" }, StringCollection should construct the figure on the right. The only difference is that instead of an EOW link, states are marked as final.

The algorithm for constructing the data structure is from the paper Incremental construction of minimal acyclic finite-state automata.

namespace Mjolka.Collections
{
    using System;
    using System.Collections.Generic;
    using System.Text;

    /// <summary>
    /// Represents a memory-efficient collection of strings.
    /// </summary>
    [Serializable]
    public class StringCollection : IEnumerable<string>
    {
        /// <summary>
        /// The initial state of the finite-state automaton.
        /// </summary>
        private readonly State initialState;

        /// <summary>
        /// The number of elements in the collection.
        /// </summary>
        private readonly int count;

        /// <summary>
        /// Initializes a new instance of the <see cref="StringCollection"/> class that contains
        /// the specified strings, where the strings are in lexicographic order.
        /// </summary>
        /// <param name="strings">
        /// The collection of strings that are in the new <see cref="StringCollection"/>, in
        /// lexicographic order.
        /// </param>
        /// <exception cref="ArgumentNullException">
        /// <paramref name="strings"/> is <c>null</c>.
        /// </exception>
        /// <exception cref="ArgumentException">
        /// An item in <paramref name="strings"/> is <c>null</c>.
        /// </exception>
        public StringCollection(IEnumerable<string> strings)
        {
            if (strings == null)
            {
                throw new ArgumentNullException("strings");
            }

            var register = new Dictionary<State, State>(new StateComparer());
            using (var enumerator = strings.GetEnumerator())
            {
                if (!enumerator.MoveNext())
                {
                    return;
                }

                if (enumerator.Current == null)
                {
                    throw new ArgumentException("strings cannot contain null", "strings");
                }

                if (string.IsNullOrEmpty(enumerator.Current))
                {
                    this.initialState = new FinalState();
                }
                else
                {
                    this.initialState = new State();
                    AddSuffix(this.initialState, enumerator.Current.ToCharArray());
                }

                this.count = 1;

                while (enumerator.MoveNext())
                {
                    var word = enumerator.Current;
                    if (word == null)
                    {
                        throw new ArgumentException("strings cannot contain null", "strings");
                    }

                    var chars = word.ToCharArray();
                    int prefix;
                    var lastState = this.CommonPrefix(chars, out prefix);
                    if (lastState.HasChildren)
                    {
                        ReplaceOrRegister(lastState, register);
                    }

                    var suffix = new ArraySegment<char>(chars, prefix, word.Length - prefix);
                    AddSuffix(lastState, suffix);
                    this.count++;
                }
            }

            ReplaceOrRegister(this.initialState, register);
            register.Clear();
        }

        /// <summary>
        /// Gets the number of strings contained in the <see cref="StringCollection"/>.
        /// </summary>
        public int Count
        {
            get { return this.count; }
        }

        /// <summary>
        /// Determines if the <see cref="StringCollection"/> contains the specified value.
        /// </summary>
        /// <param name="value">
        /// The value to locate in the <see cref="StringCollection"/>.
        /// </param>
        /// <returns>
        /// <code>true</code> if the <see cref="StringCollection"/> contains the specified value;
        /// otherwise, <code>false</code>.
        /// </returns>
        public bool Contains(string value)
        {
            if (value == null)
            {
                return false;
            }

            var state = this.initialState;
            foreach (var c in value.ToCharArray())
            {
                var nextState = state.Transition(c);
                if (nextState == null)
                {
                    return false;
                }

                state = nextState;
            }

            return state.IsFinal;
        }

        /// <summary>
        /// The get enumerator.
        /// </summary>
        /// <returns>
        /// The <see cref="Enumerator"/>.
        /// </returns>
        public Enumerator GetEnumerator()
        {
            return new Enumerator(this);
        }

        /// <summary>
        /// The get enumerator.
        /// </summary>
        /// <returns>
        /// The enumerator.
        /// </returns>
        IEnumerator<string> IEnumerable<string>.GetEnumerator()
        {
            return new Enumerator(this);
        }

        /// <summary>
        /// The get enumerator.
        /// </summary>
        /// <returns>
        /// The enumerator.
        /// </returns>
        System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator()
        {
            return new Enumerator(this);
        }

        /// <summary>
        /// Counts the number of states in the finite-state automaton.
        /// </summary>
        /// <returns>
        /// The number of states in the finite-state automaton.
        /// </returns>
        internal int CountStates()
        {
            var seen = new HashSet<State>();
            var stack = new Stack<State>();
            stack.Push(this.initialState);
            while (stack.Count > 0)
            {
                var state = stack.Pop();
                if (!seen.Add(state))
                {
                    continue;
                }

                foreach (var child in state.Children)
                {
                    stack.Push(child.State);
                }
            }

            return seen.Count;
        }

        /// <summary>
        /// Adds a suffix to a state.
        /// </summary>
        /// <param name="state">
        /// The state to add the suffix to.
        /// </param>
        /// <param name="suffix">
        /// The suffix to add to the state.
        /// </param>
        private static void AddSuffix(State state, IList<char> suffix)
        {
            var length = suffix.Count - 1;
            for (var i = 0; i < length; i++)
            {
                var newState = new State();
                state.AddEdge(suffix[i], newState);
                state = newState;
            }

            state.AddEdge(suffix[length], new FinalState());
        }

        /// <summary>
        /// Replaces a state or inserts it into the register.
        /// </summary>
        /// <param name="state">
        /// The state to replace or insert into the register.
        /// </param>
        /// <param name="register">
        /// The register.
        /// </param>
        private static void ReplaceOrRegister(State state, IDictionary<State, State> register)
        {
            var child = state.LastChild;
            if (child == null)
            {
                return;
            }

            if (child.HasChildren)
            {
                ReplaceOrRegister(child, register);
            }

            State replacement;
            if (register.TryGetValue(child, out replacement))
            {
                state.LastChild = replacement;
            }
            else
            {
                register.Add(child, child);
            }
        }

        /// <summary>
        /// Gets the common prefix of a word and the finite-state automaton.
        /// </summary>
        /// <param name="word">
        /// The word to find the common prefix of.
        /// </param>
        /// <param name="prefixLength">
        /// The length of the common prefix.
        /// </param>
        /// <returns>
        /// The last <see cref="State"/> in the common prefix.
        /// </returns>
        private State CommonPrefix(IList<char> word, out int prefixLength)
        {
            var state = this.initialState;
            var length = word.Count;
            for (var i = 0; i < length; i++)
            {
                var next = state.Transition(word[i]);
                if (next == null)
                {
                    prefixLength = i;
                    return state;
                }

                state = next;
            }

            prefixLength = length;
            return state;
        }

        /// <summary>
        /// Represents a labeled edge from one state to another.
        /// </summary>
        [Serializable]
        private struct Edge
        {
            /// <summary>
            /// The edge label.
            /// </summary>
            public readonly char Label;

            /// <summary>
            /// The edge state.
            /// </summary>
            public readonly State State;

            /// <summary>
            /// Initializes a new instance of the <see cref="Edge"/> struct.
            /// </summary>
            /// <param name="label">
            /// The edge label.
            /// </param>
            /// <param name="state">
            /// The edge state.
            /// </param>
            public Edge(char label, State state)
            {
                this.Label = label;
                this.State = state;
            }
        }

        /// <summary>
        /// The path segment.
        /// </summary>
        private struct PathSegment
        {
            /// <summary>
            /// The label.
            /// </summary>
            public readonly char Label;

            /// <summary>
            /// The depth.
            /// </summary>
            public readonly int Depth;

            /// <summary>
            /// Initializes a new instance of the <see cref="PathSegment"/> struct.
            /// </summary>
            /// <param name="label">
            /// The label.
            /// </param>
            /// <param name="depth">
            /// The depth.
            /// </param>
            public PathSegment(char label, int depth)
            {
                this.Label = label;
                this.Depth = depth;
            }
        }

        /// <summary>
        /// The enumerator.
        /// </summary>
        public class Enumerator : IEnumerator<string>
        {
            /// <summary>
            /// The initial state.
            /// </summary>
            private readonly State initialState;

            /// <summary>
            /// The string builder.
            /// </summary>
            private readonly StringBuilder stringBuilder;

            /// <summary>
            /// The states to traverse next.
            /// </summary>
            private readonly Stack<State> states;

            /// <summary>
            /// The edges leading to the states on the stack.
            /// </summary>
            private readonly Stack<PathSegment> edges;

            /// <summary>
            /// Initializes a new instance of the <see cref="Enumerator"/> class.
            /// </summary>
            /// <param name="stringCollection">
            /// The string collection.
            /// </param>
            public Enumerator(StringCollection stringCollection)
            {
                this.initialState = stringCollection.initialState;
                this.stringBuilder = new StringBuilder();
                this.states = new Stack<State>();
                this.states.Push(this.initialState);
                this.edges = new Stack<PathSegment>();
            }

            /// <inheritdoc />
            public string Current
            {
                get { return this.stringBuilder.ToString(); }
            }

            /// <inheritdoc />
            object System.Collections.IEnumerator.Current
            {
                get { return this.Current; }
            }

            /// <inheritdoc />
            public bool MoveNext()
            {
                // Move to the next FinalState, adding chars in the
                // path to this.stringBuilder.
                while (this.states.Count > 0)
                {
                    var state = this.states.Pop();
                    var depth = -1;
                    if (this.edges.Count > 0)
                    {
                        var edge = this.edges.Pop();
                        depth = edge.Depth;
                        if (depth < this.stringBuilder.Length)
                        {
                            this.stringBuilder.Remove(depth, this.stringBuilder.Length - depth);
                        }

                        var label = edge.Label;
                        this.stringBuilder.Append(label);
                    }

                    var children = state.Children;
                    if (children != null)
                    {
                        depth = depth + 1;
                        for (var i = children.Length - 1; i >= 0; i--)
                        {
                            var child = children[i];
                            this.edges.Push(new PathSegment(child.Label, depth));
                            this.states.Push(child.State);
                        }
                    }

                    if (state.IsFinal)
                    {
                        return true;
                    }
                }

                return false;
            }

            /// <inheritdoc />
            public void Reset()
            {
                this.stringBuilder.Clear();
                this.states.Clear();
                this.states.Push(this.initialState);
                this.edges.Clear();
            }

            /// <inheritdoc />
            public void Dispose()
            {
            }
        }

        /// <summary>
        /// Represents a state in the finite-state automaton.
        /// </summary>
        [Serializable]
        private class State
        {
            /// <summary>
            /// The edges out of the state.
            /// </summary>
            private Edge[] children = new Edge[0];

            /// <summary>
            /// Gets a value indicating whether the state has children.
            /// </summary>
            public bool HasChildren
            {
                get { return this.children.Length > 0; }
            }

            /// <summary>
            /// Gets or sets the last edge out of the state.
            /// </summary>
            public State LastChild
            {
                get
                {
                    var length = this.children.Length;
                    return length > 0
                        ? this.children[length - 1].State
                        : null;
                }

                set
                {
                    var lastIndex = this.children.Length - 1;
                    var lastChild = this.children[lastIndex];
                    this.children[lastIndex] = new Edge(lastChild.Label, value);
                }
            }

            /// <summary>
            /// Gets a value indicating whether the state is final.
            /// </summary>
            public virtual bool IsFinal
            {
                get { return false; }
            }

            /// <summary>
            /// Gets the children.
            /// </summary>
            public Edge[] Children
            {
                get { return this.children; }
            }

            /// <summary>
            /// Gets the state corresponding to the edge with the specified label out of the state,
            /// if it exists.
            /// </summary>
            /// <param name="label">
            /// The label of the edge out of the state.
            /// </param>
            /// <returns>
            /// The <see cref="State"/> corresponding to the edge with the specified label.
            /// </returns>
            public State Transition(char label)
            {
                var length = this.children.Length;
                for (var i = 0; i < length; i++)
                {
                    var child = this.children[i];
                    if (child.Label == label)
                    {
                        return child.State;
                    }
                }

                return null;
            }

            /// <summary>
            /// Adds an edge out of the state.
            /// </summary>
            /// <param name="label">
            /// The label of the edge to the specified state.
            /// </param>
            /// <param name="state">
            /// The state.
            /// </param>
            public void AddEdge(char label, State state)
            {
                var length = this.children.Length;
                var newChildren = new Edge[length + 1];

                Array.Copy(this.children, newChildren, length);

                newChildren[length] = new Edge(label, state);
                this.children = newChildren;
            }
        }

        /// <summary>
        /// Represents a final state in the finite-state automaton.
        /// </summary>
        [Serializable]
        private sealed class FinalState : State
        {
            /// <inheritdoc />
            public override bool IsFinal
            {
                get { return true; }
            }
        }

        /// <summary>
        /// Represents a state comparer.
        /// </summary>
        private class StateComparer : IEqualityComparer<State>
        {
            /// <inheritdoc />
            public bool Equals(State left, State right)
            {
                if (left.IsFinal != right.IsFinal)
                {
                    return false;
                }

                var leftChildren = left.Children;
                var rightChildren = right.Children;

                if (leftChildren.Length != rightChildren.Length)
                {
                    return false;
                }

                for (var i = 0; i < leftChildren.Length; i++)
                {
                    var leftChild = leftChildren[i];
                    var rightChild = rightChildren[i];
                    if (leftChild.Label != rightChild.Label || leftChild.State != rightChild.State)
                    {
                        return false;
                    }
                }

                return true;
            }

            /// <inheritdoc />
            public int GetHashCode(State obj)
            {
                unchecked
                {
                    var result = obj.IsFinal ? 17 : 521;
                    foreach (var child in obj.Children)
                    {
                        result = (31 * result) + child.Label.GetHashCode();
                        result = (31 * result) + child.State.GetHashCode();
                    }

                    return result;
                }
            }
        }
    }
}

The code and unit tests are available on GitHub.

All feedback/improvements are welcome. There are some parts I'm not particularly happy with:

• The special-casing of the empty string in the constructor
• Using an out parameter in CommonPrefix
• All of Enumerator.MoveNext

Edit: Subclassing State and using Edge[] instead of List<Edge> were guided by profiling. Here are some stats from dotMemory:

Test input: BritishWords.txt. 79,764 words, 776,045 bytes.

Build: .NET 4.5/x86/Release.

Code as posted here:

                                          Retained Bytes
StringCollection                               1,469,008

Using a boolean field for IsFinal instead of subclassing State:

StringCollection                               1,612,224

Using List<Edge> (andTrimExcess) instead of Edge[]:

StringCollection                               2,360,732

In comparison, using a HashSet<string> (which requires the strings to also be in memory):

                            Objects Count          Bytes
string                             79,764      2,509,260

                                          Retained Bytes
HashSet<string>                         1      1,448,432
                                               ---------
                                               3,957,692

Of course, for many applications, 2.5MB is just a rounding error, but sometimes it matters. I'm just trying to reduce memory usage as a learning exercise.

Answer 1

For this answer, I'll work on the following assumptions, so tell me if any of them are wrong:

• You're most interested in memory efficiency
• You're not so interested in time efficiency, especially within the initial construction of the collection
• You're happy to use LINQ

Enumerator

The first thing I notice here are the stacks. And the first thing that makes me think: recursion! You're not only worrying about keeping two stacks aligned, but you also have to keep track of how deep you are inside them. Taking advantage of the call stack to take care of that for you makes life a lot easier.

So start with the signature:

private IEnumerable<string> Next(State state, PathSegment? edge, StringBuilder builder)

You know the State and PathSegment? need to be in there because they're the ones in stacks. The way I've done it I've replaced the class with an iterator method, which necessitates the StringBuilder being passed back and forth rather than being kept as a class-level variable.

From there, it's fairly simple to rework what you have currently to a new version:

private IEnumerable<string> Next(State state, PathSegment? edge, StringBuilder builder)
{
    if (edge != null)
    {
        builder.Append(edge.Value.Label);
    }

    if (state.IsFinal)
        yield return builder.ToString();

    foreach (var child in state.Children)
    {
        var nextEdge = new PathSegment(child.Label, 0);
        foreach (var result in Next(child.State, nextEdge, builder))
            yield return result;
    }
    if(builder.Length > 0)
        builder.Remove(builder.Length - 1, 1);
}

Notice that by taking advantage of the call stack, depth isn't needed at all, we just remove the last character from the builder before exiting the method. This means that we can actually get rid of the PathSegment class altogether and use Edge directly instead!

After applying that, the next thing that stands out is this line:

foreach (var result in Next(child.State, child, builder))

It doesn't seem like we should have to pass both child.State and child. So let's simplify further:

private IEnumerable<string> Next(Edge edge, StringBuilder builder)
{
    var state = edge.State;
    if (state.IsFinal)
        yield return builder.ToString();

    foreach (var child in state.Children)
    {
        builder.Append(child.Label);
        foreach (var result in Next(child, builder))
            yield return result;
    }
    if (builder.Length > 0)
        builder.Remove(builder.Length - 1, 1);
}

Note that labels are now appended before calling into the method recursively, rather than at the beginning. This is so that we can do:

public IEnumerable<string> StringEnumerator()
{
    var entryEdge = new Edge('\0', initialState);
    return Next(entryEdge, new StringBuilder());
}

You can plonk this directly in your StringCollection class and use it for your GetEnumerator.

Answer 2

StringCollection

Naming

StringCollection is, as I pointed out in the comments, a misleading name that suggests your type is implementing ICollection<T>. A better name is pretty hard to find though, since EnumerableString just looks weird, as String already implements IEnumerable<char>.

The only "better" name I can think of, is Strings... it's somewhat lame, granted.. but it's not misleading.

The Contains member method, because of how the C# compiler resolves method calls, will always be called over Linq's Contains<TSource>(IEnumerable<TSource>, TSource) extension method; this may or may not be desired, and/or expected. Since you already have XML comments, you could add a <remarks> tag, and document that.

---

Duplication

You return new Enumerator(this) a lot. Why not have a private IEnumerator GetEnumeratorInternal() method that all GetEnumerator() methods could call?

---

Constructor

I wouldn't bother with the early return if the MoveNext returns false, since the while loop won't iterate in that case, and you'll no longer need to check whether Current is null outside the loop, and you can start incrementing the count within the loop instead of starting at 1.

---

Enumerator

this

This is purely aesthetic, but I find the constructor... bulky:

public Enumerator(StringCollection stringCollection)
{
    this.initialState = stringCollection.initialState;
    this.stringBuilder = new StringBuilder();
    this.states = new Stack<State>();
    this.states.Push(this.initialState);
    this.edges = new Stack<PathSegment>();
}

...compared to how I would have written it:

public Enumerator(StringCollection stringCollection)
{
    _initialState = stringCollection.initialState;

    _stringBuilder = new StringBuilder();
    _edges = new Stack<PathSegment>();
    _states = new Stack<State>();

    _states.Push(this.initialState);
}

By prefixing private fields with an underscore, you remove the need to qualify them in assignations like this.

---

Comments

These comments are redundant and useless. Remove them. Don't write XML comments just for the sake of writing XML comments; private members don't/shouldn't need XML comments:

/// <summary>
/// The initial state.
/// </summary>
private readonly State initialState;

/// <summary>
/// The string builder.
/// </summary>
private readonly StringBuilder stringBuilder;

Compare to:

private readonly State _initialState;
private readonly StringBuilder _stringBuilder;

XML comments that add no value, are just noise and clutter. Burn.

---

You're declaring a very short-lived variable in MoveNext, that can afford to disappear:

var label = edge.Label;
this.stringBuilder.Append(label);

this.stringBuilder.Append(edge.Label); is clear enough IMO.

Answer 3

The biggest problem I found is that your code doesn't work when you have a shorter string (like "a") after a longer string (like "abc").

It's possible there are other cases where your code doesn't work, but I wasn't able to follow the logic of your code properly.

---

string.IsNullOrEmpty(enumerator.Current)

You already tested Current for null, so now you can only test it for == string.Empty.

---

using (var enumerator = strings.GetEnumerator())

This is ugly. I would prefer a normal foreach instead, though it would probably require changes in other methods.

---

register.Clear();

register is a local variable, that will become eligible for GC after the method exits. There is no reason to manually Clear() it just before that.

---

foreach (var c in value.ToCharArray())

string is enumerable, so you write just:

foreach (var c in value)

---

public Enumerator GetEnumerator()
IEnumerator<string> IEnumerable<string>.GetEnumerator()

This is usually done when Enumerator is a struct, which can make foreach more efficient (no heap allocations). But your Enumerator is a class, so doing this won't help you much.

---

public void AddEdge(char label, State state)

All that array copying might result in less used memory after the state machine is built, but it makes the building itself more memory-intensive and less efficient.

A better solution would be to use List here and then call TrimExcess() on all lists after you're done building.

---

private sealed class FinalState : State
{
    /// <inheritdoc />
    public override bool IsFinal
    {
        get { return true; }
    }
}

This doesn't seem to be worth creating a separate type for.

---

private class StateComparer : IEqualityComparer<State>

This seems to be the only reasonable way of comparing states, so I would make State implement IEquatable<State> instead.