Imports System.Threading
Imports System.Collections.ObjectModel
Imports Mercury.Mobile.Threading
Imports Mercury.Mobile

Namespace Collections.Concurrent

    ''' <summary>
    ''' A thread-safe collection of key-value pairs that can be accessed by multiple threads concurrently. 
    ''' </summary>
    ''' <typeparam name="TKey">The type of the key.</typeparam>
    ''' <typeparam name="TValue">The type of the value.</typeparam>
    ''' <remarks></remarks>
    Public Class ConcurrentDictionary(Of TKey, TValue)
        Implements IEnumerable(Of KeyValuePair(Of TKey, TValue))

        Friend Shared ReadOnly Property DefaultConcurrency() As Integer
            Get
                Return 4
            End Get
        End Property

        Private Class Node
            Private m_next As Node
            Friend Key As TKey
            Friend Value As TValue
            Friend Hashcode As Integer

            Public Sub New()
            End Sub

            Public Sub New(ByVal key As TKey, ByVal value As TValue, ByVal hashcode As Integer, ByVal nxt As Node)
                Me.Key = key
                Me.Value = value
                Me.Hashcode = hashcode
                Me.SetNextNode(nxt)
            End Sub

            Friend ReadOnly Property NextNode() As Node
                Get
                    Return Volatile.VolatileRead(Me.m_next)
                End Get
            End Property

            Friend Sub SetNextNode(ByVal value As Node)
                Volatile.VolatileWrite(Me.m_next, value)
            End Sub
        End Class

        Private m_buckets() As Node
        Private m_locks() As Object
        Private m_countPerLock() As Integer
        Private comparer As IEqualityComparer(Of TKey)

        ''' <summary>
        ''' Initializes a default instance of the <see cref="Mercury.Mobile.Collections.Concurrent.ConcurrentDictionary(Of TKey, TValue)"/> class.
        ''' </summary>
        ''' <remarks></remarks>
        Public Sub New()
            Me.New(ConcurrentDictionary(Of TKey, TValue).DefaultConcurrency, 31)
        End Sub

        ''' <summary>
        ''' Initializes a default instance of the <see cref="Mercury.Mobile.Collections.Concurrent.ConcurrentDictionary(Of TKey, TValue)"/> 
        ''' class with the specified concurrency level and initial capacity.
        ''' </summary>
        ''' <param name="concurrencyLevel">The estimated number of threads that will update this collection concurrently.</param>
        ''' <param name="capacity">The initial number of elements this collection can contain.</param>
        ''' <remarks></remarks>
        Public Sub New(ByVal concurrencyLevel As Integer, ByVal capacity As Integer)
            Me.New(concurrencyLevel, capacity, EqualityComparer(Of TKey).Default)
        End Sub

        ''' <summary>
        ''' Initializes a default instance of the <see cref="Mercury.Mobile.Collections.Concurrent.ConcurrentDictionary(Of TKey, TValue)"/> 
        ''' class with the initial collection of key-value pairs, and uses the specified 
        ''' <see cref="System.Collections.Generic.IEqualityComparer(of T)"/>.
        ''' </summary>
        ''' <param name="collection">A collection of initial key-value pairs to be stored in this collection.</param>
        ''' <param name="comparer">The <see cref="System.Collections.Generic.IEqualityComparer(of T)"/> 
        ''' implementation to use when comparing keys.</param>
        ''' <remarks></remarks>
        Public Sub New(ByVal collection As IEnumerable(Of KeyValuePair(Of TKey, TValue)), ByVal comparer As IEqualityComparer(Of TKey))
            Me.New(ConcurrentDictionary(Of TKey, TValue).DefaultConcurrency, collection, comparer)
        End Sub

        ''' <summary>
        ''' Initializes a default instance of the <see cref="Mercury.Mobile.Collections.Concurrent.ConcurrentDictionary(Of TKey, TValue)"/> 
        ''' class with the specified concurrency level, has the initial collection of key-value pairs, 
        ''' and uses the specified <see cref="System.Collections.Generic.IEqualityComparer(of T)"/>.
        ''' </summary>
        ''' <param name="concurrencyLevel">The estimated number of threads that will update this collection concurrently.</param>
        ''' <param name="collection">A collection of initial key-value pairs to be stored in this collection.</param>
        ''' <param name="comparer">The <see cref="System.Collections.Generic.IEqualityComparer(of T)"/> 
        ''' implementation to use when comparing keys.</param>
        ''' <remarks></remarks>
        Public Sub New(ByVal concurrencyLevel As Integer, ByVal collection As IEnumerable(Of KeyValuePair(Of TKey, TValue)), ByVal comparer As IEqualityComparer(Of TKey))
            Me.New(concurrencyLevel, 31, comparer)
            If collection Is Nothing Then _
                Throw New ArgumentNullException("collection")
            If comparer Is Nothing Then _
                Throw New ArgumentNullException("comparer")
            Me.InitializeFromCollection(collection)
        End Sub

        ''' <summary>
        ''' Initializes a default instance of the <see cref="Mercury.Mobile.Collections.Concurrent.ConcurrentDictionary(Of TKey, TValue)"/> 
        ''' class with the specified concurrency level, initial capacity, and uses the 
        ''' specified <see cref="System.Collections.Generic.IEqualityComparer(of T)"/>.
        ''' </summary>
        ''' <param name="concurrencyLevel">The estimated number of threads that will update this collection concurrently.</param>
        ''' <param name="capacity">The initial number of elements this collection can contain.</param>
        ''' <param name="comparer">The <see cref="System.Collections.Generic.IEqualityComparer(of T)"/> 
        ''' implementation to use when comparing keys.</param>
        ''' <remarks></remarks>
        Public Sub New(ByVal concurrencyLevel As Integer, ByVal capacity As Integer, ByVal comparer As IEqualityComparer(Of TKey))
            If concurrencyLevel < 1 Then _
                Throw New ArgumentOutOfRangeException("concurrencyLevel", "Concurrency level must be greater than or equal to 1.")
            If capacity < 0 Then _
                Throw New ArgumentOutOfRangeException("capacity", "Capacity must be greater than or equal to 0.")
            If comparer Is Nothing Then _
                Throw New ArgumentNullException("comparer")
            If capacity < concurrencyLevel Then _
                capacity = concurrencyLevel
            Me.m_locks = New Object(concurrencyLevel - 1) {}
            For i As Integer = 0 To Me.m_locks.Length - 1
                Me.m_locks(i) = New Object()
            Next
            Me.m_countPerLock = New Integer(Me.m_locks.Length - 1) {}
            Me.m_buckets = New Node(capacity - 1) {}
            Me.comparer = comparer
        End Sub

        Private Sub InitializeFromCollection(ByVal collection As IEnumerable(Of KeyValuePair(Of TKey, TValue)))
            For Each current As KeyValuePair(Of TKey, TValue) In collection
                If current.Key Is Nothing Then _
                    Throw New ArgumentNullException("key")
                Dim value As TValue = Nothing
                If Not (Me.TryAddInternal(current.Key, current.Value, False, False, value)) Then
                    Throw New ArgumentException("The source collection contains duplicate keys.")
                End If
            Next
        End Sub

        ''' <summary>
        ''' Removes all keys and values from the <see cref="Mercury.Mobile.Collections.Concurrent.ConcurrentDictionary(Of TKey, TValue)"/>.
        ''' </summary>
        ''' <remarks></remarks>
        Public Sub Clear()
            Dim toExclusive As Integer = 0
            Try
                Me.AcquireAllLocks(toExclusive)
                Me.m_buckets = New Node(30) {}
                Array.Clear(Me.m_countPerLock, 0, Me.m_countPerLock.Length)
            Finally
                Me.ReleaseLocks(0, toExclusive)
            End Try
        End Sub

        ''' <summary>
        ''' Determines whether the key exists in this <see cref="Mercury.Mobile.Collections.Concurrent.ConcurrentDictionary(Of K, V)"/>.
        ''' </summary>
        ''' <param name="key">The key to search for.</param>
        ''' <returns>True if the key exists in this <see cref="Mercury.Mobile.Collections.Concurrent.ConcurrentDictionary(Of K, V)"/>; otherwise, false.</returns>
        ''' <remarks></remarks>
        Public Function ContainsKey(ByVal key As TKey) As Boolean
            If key Is Nothing Then _
                Throw New ArgumentNullException("key")
            Dim value As TValue = Nothing
            Return Me.TryGetValue(key, value)
        End Function

        ''' <summary>
        ''' Returns an enumerator that iterates through the collection of key-value pairs.
        ''' </summary>
        ''' <returns>An enumerator that iterates through the collection of key-value pairs.</returns>
        ''' <remarks></remarks>
        Public Function GetEnumerator() As System.Collections.Generic.IEnumerator(Of System.Collections.Generic.KeyValuePair(Of TKey, TValue)) Implements System.Collections.Generic.IEnumerable(Of System.Collections.Generic.KeyValuePair(Of TKey, TValue)).GetEnumerator
            Dim initial = New Tuple(Of Node(), Node, Integer)(Me.m_buckets, Nothing, -1)
            Return New Iterator(Of KeyValuePair(Of TKey, TValue), Tuple(Of Node(), Node, Integer))(AddressOf __MoveNext, initial)
        End Function

        Private Function __MoveNext(ByRef nextItem As KeyValuePair(Of TKey, TValue), ByRef state As Tuple(Of Node(), Node, Integer)) As Boolean
            Dim buckets = state.Item1
            Dim node = state.Item2
            Dim idx = state.Item3
            If node Is Nothing Then
                idx += 1
                If idx >= buckets.Length Then _
                    Return False
                node = buckets(idx)
                Thread.MemoryBarrier()
            End If
            If node IsNot Nothing Then
                nextItem = New KeyValuePair(Of TKey, TValue)(node.Key, node.Value)
                state = New Tuple(Of Node(), Node, Integer)(buckets, node.NextNode, idx)
                Return True
            End If
            Return False
        End Function

        Private Function GetEnumerator1() As System.Collections.IEnumerator Implements System.Collections.IEnumerable.GetEnumerator
            Return Me.GetEnumerator()
        End Function

        ''' <summary>
        ''' Gets the number of elements in this dictionary.
        ''' </summary>
        ''' <returns>The number of elements in this dictionary.</returns>
        ''' <remarks></remarks>
        Public ReadOnly Property Count() As Integer
            Get
                Dim num As Integer = 0
                Dim toExclusive As Integer = 0
                Try
                    Me.AcquireAllLocks(toExclusive)
                    For i As Integer = 0 To Me.m_countPerLock.Length - 1
                        num += Me.m_countPerLock(i)
                    Next
                Finally
                    Me.ReleaseLocks(0, toExclusive)
                End Try
                Return num
            End Get
        End Property

        ''' <summary>
        ''' Gets an <see cref="System.Collections.ICollection"/> containing the keys in this dictionary.
        ''' </summary>
        ''' <returns>An <see cref="System.Collections.ICollection"/> containing the keys in this dictionary.</returns>
        ''' <remarks></remarks>
        Public ReadOnly Property Keys() As ICollection(Of TKey)
            Get
                Return Me.GetKeys()
            End Get
        End Property

        ''' <summary>
        ''' Gets an <see cref="System.Collections.ICollection"/> containing the values in this dictionary.
        ''' </summary>
        ''' <returns>An <see cref="System.Collections.ICollection"/> containing the values in this dictionary.</returns>
        ''' <remarks></remarks>
        Public ReadOnly Property Values() As ICollection(Of TValue)
            Get
                Return Me.GetValues()
            End Get
        End Property

        ''' <summary>
        ''' Gets or sets the value associated with the specified key.
        ''' </summary>
        ''' <param name="key">The key of the associated value to get or set.</param>
        ''' <returns>The value associated with the specified key.</returns>
        ''' <remarks></remarks>
        Public Property Item(ByVal key As TKey) As TValue
            Get
                Dim v As TValue
                If Not (Me.TryGetValue(key, v)) Then _
                    Throw New KeyNotFoundException()
                Return v
            End Get
            Set(ByVal value As TValue)
                If key Is Nothing Then _
                    Throw New ArgumentNullException("key")
                Dim tVal As TValue = Nothing
                Me.TryAddInternal(key, value, True, True, tVal)
            End Set
        End Property

        ''' <summary>
        ''' Attempts to add the specified key and value to the dictionary.
        ''' </summary>
        ''' <param name="key">The key to add.</param>
        ''' <param name="value">The associated value to add.</param>
        ''' <returns>True if the key/value pair was added to the dictionary successfully; otherwise, false.</returns>
        ''' <remarks></remarks>
        Public Function TryAdd(ByVal key As TKey, ByVal value As TValue) As Boolean
            If key Is Nothing Then _
                Throw New ArgumentNullException("key")
            Dim tVal As TValue = Nothing
            Return Me.TryAddInternal(key, value, False, True, tVal)
        End Function

        Private Function TryAddInternal(ByVal key As TKey, ByVal value As TValue, ByVal updateIfExists As Boolean, ByVal acquireLock As Boolean, ByRef resultValue As TValue) As Boolean
            Dim hashcode = Me.comparer.GetHashCode(key)
            Dim buckets() As ConcurrentDictionary(Of TKey, TValue).Node = Nothing
            Dim flag As Boolean
            Do While True
                buckets = Me.m_buckets
                Dim bucketNo As Integer
                Dim lockNo As Integer
                Me.GetBucketAndLockNo(hashcode, bucketNo, lockNo, buckets.Length)
                flag = False
                If acquireLock Then _
                    Monitor.Enter(Me.m_locks(lockNo))
                Try
                    If Not (Object.ReferenceEquals(buckets, Me.m_buckets)) Then _
                        Continue Do
                    Dim node As ConcurrentDictionary(Of TKey, TValue).Node = Nothing
                    Dim n2 = buckets(bucketNo)
                    Do While n2 IsNot Nothing
                        If Me.comparer.Equals(n2.Key, key) Then
                            If updateIfExists Then
                                Dim node3 = New Node(n2.Key, n2.Value, hashcode, n2.NextNode)
                                If node Is Nothing Then
                                    buckets(bucketNo) = node3
                                Else
                                    node.SetNextNode(node3)
                                End If
                                resultValue = value
                            Else
                                resultValue = n2.Value
                            End If
                            Return False
                        End If
                        node = n2
                        n2 = n2.NextNode
                    Loop
                    buckets(bucketNo) = New Node(key, value, hashcode, buckets(bucketNo))
                    Interlocked.Increment(Me.m_countPerLock(lockNo))
                    If Volatile.VolatileRead(Me.m_countPerLock(lockNo)) > (buckets.Length / Me.m_locks.Length) Then
                        flag = True
                    End If
                Finally
                    If acquireLock Then _
                        Monitor.Exit(Me.m_locks(lockNo))
                End Try
                Exit Do
            Loop
            If flag Then _
                Me.GrowTable(buckets)
            resultValue = value
            Return True
        End Function

        ''' <summary>
        ''' Attempts to get the value associated with the specified key.
        ''' </summary>
        ''' <param name="key">The key of the associated value to search for.</param>
        ''' <param name="value">Out. The found value.</param>
        ''' <returns>True if the key was found in the dictionary; otherwise, false.</returns>
        ''' <remarks></remarks>
        Public Function TryGetValue(ByVal key As TKey, ByRef value As TValue) As Boolean
            If key Is Nothing Then _
                Throw New ArgumentNullException("key")
            Dim buckets = Me.m_buckets
            Dim bucketNo As Integer
            Dim lockNo As Integer
            Me.GetBucketAndLockNo(Me.comparer.GetHashCode(key), bucketNo, lockNo, buckets.Length)
            Dim n = buckets(bucketNo)
            Thread.MemoryBarrier()
            Do While n IsNot Nothing
                If Me.comparer.Equals(n.Key, key) Then
                    value = n.Value
                    Return True
                End If
                n = n.NextNode
            Loop
            value = Nothing
            Return False
        End Function

        ''' <summary>
        ''' Attempts to remove and return the value that has the specified key.
        ''' </summary>
        ''' <param name="key">The key of the element to remove.</param>
        ''' <param name="value">Out. The value associated with the specified 
        ''' key or the default value if the key does not exist.</param>
        ''' <returns>True if the object was successfully removed; otherwise, false.</returns>
        ''' <remarks></remarks>
        Public Function TryRemove(ByVal key As TKey, ByRef value As TValue) As Boolean
            If key Is Nothing Then _
                Throw New ArgumentNullException("key")
            Return Me.TryRemoveInternal(key, value)
        End Function

        Private Function TryRemoveInternal(ByVal key As TKey, ByRef value As TValue) As Boolean
            Do While True
                Dim buckets = Me.m_buckets
                Dim bucketNo As Integer
                Dim lockNo As Integer
                Me.GetBucketAndLockNo(Me.comparer.GetHashCode(key), bucketNo, lockNo, Me.m_buckets.Length)
                SyncLock Me.m_locks(lockNo)
                    If Not (Object.ReferenceEquals(buckets, Me.m_buckets)) Then _
                        Continue Do
                    Dim nPrev As Node = Nothing
                    Dim nCurr As Node = Me.m_buckets(bucketNo)
                    Do While nCurr IsNot Nothing
                        If Me.comparer.Equals(nCurr.Key, key) Then
                            If nPrev Is Nothing Then
                                Me.m_buckets(bucketNo) = nCurr.NextNode
                            Else
                                nPrev.SetNextNode(nCurr.NextNode)
                            End If
                            value = nCurr.Value
                            Me.m_countPerLock(lockNo) -= 1
                            Return True
                        Else
                            nPrev = nCurr
                            nCurr = nPrev.NextNode
                        End If
                    Loop
                End SyncLock
                Exit Do
            Loop
            value = Nothing
            Return False
        End Function

        Private Function GetKeys() As ReadOnlyCollection(Of TKey)
            Dim toExclusive = 0
            Dim result As ReadOnlyCollection(Of TKey)
            Try
                Me.AcquireAllLocks(toExclusive)
                Dim list = New List(Of TKey)()
                For i As Integer = 0 To Me.m_buckets.Length - 1
                    Dim node As Node = Me.m_buckets(i)
                    Do While node IsNot Nothing
                        list.Add(node.Key)
                        node = node.NextNode
                    Loop
                Next
                result = New ReadOnlyCollection(Of TKey)(list)
            Finally
                Me.ReleaseLocks(0, toExclusive)
            End Try
            Return result
        End Function

        Private Function GetValues() As ReadOnlyCollection(Of TValue)
            Dim toExclusive = 0
            Dim result As ReadOnlyCollection(Of TValue)
            Try
                Me.AcquireAllLocks(toExclusive)
                Dim list = New List(Of TValue)()
                For i As Integer = 0 To Me.m_buckets.Length - 1
                    Dim node As Node = Me.m_buckets(i)
                    Do While node IsNot Nothing
                        list.Add(node.Value)
                        node = node.NextNode
                    Loop
                Next
                result = New ReadOnlyCollection(Of TValue)(list)
            Finally
                Me.ReleaseLocks(0, toExclusive)
            End Try
            Return result
        End Function

        Private Sub GrowTable(ByVal buckets() As Node)
            Dim toExclusive As Integer = 0
            Try
                Me.AcquireLocks(0, 1, toExclusive)
                If Object.ReferenceEquals(buckets, Me.m_buckets) Then
                    Dim arr() As Node
                    Dim arr2() As Integer
                    Dim num As Integer
                    Try
                        num = buckets.Length * 2 + 1
                        Do While (num Mod 3) = 0 OrElse (num Mod 5) = 0 OrElse (num Mod 7) = 0
                            num += 2
                        Loop
                    Catch ex As OverflowException
                        Return
                    End Try
                    arr = New Node(num - 1) {}
                    arr2 = New Integer(Me.m_locks.Length - 1) {}
                    Me.AcquireLocks(1, Me.m_locks.Length, toExclusive)

                    For i As Integer = 0 To buckets.Length - 1
                        Dim node As Node = buckets(i)
                        Dim nxt As Node = Nothing
                        Do While node IsNot Nothing
                            nxt = node.NextNode
                            Dim bucketNo As Integer
                            Dim lockNo As Integer
                            Me.GetBucketAndLockNo(node.Hashcode, bucketNo, lockNo, arr.Length)
                            arr(bucketNo) = New Node(node.Key, node.Value, node.Hashcode, arr(bucketNo))
                            arr2(lockNo) += 1
                            node = nxt
                        Loop
                    Next
                    Me.m_buckets = arr
                    Me.m_countPerLock = arr2
                End If
            Finally
                Me.ReleaseLocks(0, toExclusive)
            End Try
        End Sub

        Private Sub AcquireAllLocks(ByRef locksAcquired As Integer)
            Me.AcquireLocks(0, Me.m_locks.Length, locksAcquired)
        End Sub

        Private Sub AcquireLocks(ByVal fromInclusive As Integer, ByVal toExclusive As Integer, ByRef locksAcquired As Integer)
            For i As Integer = fromInclusive To toExclusive - 1
                Dim flag = False
                Try
                    Monitor.Enter(Me.m_locks(i))
                    flag = True
                Finally
                    If flag Then _
                        locksAcquired += 1
                End Try
            Next
        End Sub

        Private Sub ReleaseLocks(ByVal fromInclusive As Integer, ByVal toExclusive As Integer)
            For i As Integer = fromInclusive To toExclusive - 1
                Monitor.Exit(Me.m_locks(i))
            Next
        End Sub

        Private Sub GetBucketAndLockNo(ByVal hashcode As Integer, ByRef bucketNo As Integer, ByRef lockNo As Integer, ByVal bucketCount As Integer)
            bucketNo = ((hashcode And Integer.MaxValue) Mod bucketCount)
            lockNo = (bucketNo Mod Me.m_locks.Length)
        End Sub

    End Class

End Namespace