Imports System.Threading
Imports System.Collections.Generic

Namespace Threading

    ''' <summary>
    ''' Represents a thread-safe and framework independent collection of key/value 
    ''' pairs that can be accessed by multiple threads concurrently.
    ''' </summary>
    ''' <remarks></remarks>
    Public Class ThreadSafeDictionary(Of TKey, TValue)
        Implements IDictionary(Of TKey, TValue)

        Public Sub New()
            Me._internalDictionary = New Dictionary(Of TKey, TValue)()
        End Sub

        Private ReadOnly m_lock As Object = New Object()

        Private _internalDictionary As Dictionary(Of TKey, TValue)

#Region "ICollection"
        Public Sub Add(ByVal item As System.Collections.Generic.KeyValuePair(Of TKey, TValue)) Implements System.Collections.Generic.ICollection(Of System.Collections.Generic.KeyValuePair(Of TKey, TValue)).Add
            Me.TryAdd(item.Key, item.Value)
        End Sub

        Public Sub Clear() Implements System.Collections.Generic.ICollection(Of System.Collections.Generic.KeyValuePair(Of TKey, TValue)).Clear
            SyncLock Me.m_lock
                Me._internalDictionary.Clear()
            End SyncLock
        End Sub

        Public Function Contains(ByVal item As System.Collections.Generic.KeyValuePair(Of TKey, TValue)) As Boolean Implements System.Collections.Generic.ICollection(Of System.Collections.Generic.KeyValuePair(Of TKey, TValue)).Contains
            Return Me._internalDictionary.Contains(item)
        End Function

        Public Sub CopyTo(ByVal array() As System.Collections.Generic.KeyValuePair(Of TKey, TValue), ByVal arrayIndex As Integer) Implements System.Collections.Generic.ICollection(Of System.Collections.Generic.KeyValuePair(Of TKey, TValue)).CopyTo
            Throw New NotImplementedException()
        End Sub

        Public ReadOnly Property Count() As Integer Implements System.Collections.Generic.ICollection(Of System.Collections.Generic.KeyValuePair(Of TKey, TValue)).Count
            Get
                Return Me._internalDictionary.Count
            End Get
        End Property

        Public ReadOnly Property IsReadOnly() As Boolean Implements System.Collections.Generic.ICollection(Of System.Collections.Generic.KeyValuePair(Of TKey, TValue)).IsReadOnly
            Get
                Return False
            End Get
        End Property

        Public Function Remove(ByVal item As System.Collections.Generic.KeyValuePair(Of TKey, TValue)) As Boolean Implements System.Collections.Generic.ICollection(Of System.Collections.Generic.KeyValuePair(Of TKey, TValue)).Remove
            Return Me.TryRemove(item.Key, item.Value)
        End Function
#End Region

#Region "IDictionary"

        Public Sub Add1(ByVal key As TKey, ByVal value As TValue) Implements System.Collections.Generic.IDictionary(Of TKey, TValue).Add
            Me.TryAdd(key, value)
        End Sub

        Public Function ContainsKey(ByVal key As TKey) As Boolean Implements System.Collections.Generic.IDictionary(Of TKey, TValue).ContainsKey
            Return Me._internalDictionary.ContainsKey(key)
        End Function

        Default Public Property Item(ByVal key As TKey) As TValue Implements System.Collections.Generic.IDictionary(Of TKey, TValue).Item
            Get
                Dim value As TValue = Nothing
                Me.TryGetValue(key, value)
                Return value
            End Get
            Set(ByVal value As TValue)
                If Me.ContainsKey(key) Then
                    SyncLock Me.m_lock
                        If Me.ContainsKey(key) Then _
                            Me._internalDictionary.Item(key) = value
                    End SyncLock
                End If
            End Set
        End Property

        Public ReadOnly Property Keys() As System.Collections.Generic.ICollection(Of TKey) Implements System.Collections.Generic.IDictionary(Of TKey, TValue).Keys
            Get
                SyncLock Me.m_lock
                    Return Me._internalDictionary.Keys
                End SyncLock
            End Get
        End Property

        Public Function Remove1(ByVal key As TKey) As Boolean Implements System.Collections.Generic.IDictionary(Of TKey, TValue).Remove
            Dim value As TValue = Nothing
            Return Me.TryRemove(key, value)
        End Function

        Public Function TryGetValue(ByVal key As TKey, ByRef value As TValue) As Boolean Implements System.Collections.Generic.IDictionary(Of TKey, TValue).TryGetValue
            If Me.ContainsKey(key) Then
                SyncLock Me.m_lock
                    If Me.ContainsKey(key) Then
                        value = Me._internalDictionary(key)
                        Return True
                    End If
                End SyncLock
            End If
            Return False
        End Function

        Public ReadOnly Property Values() As System.Collections.Generic.ICollection(Of TValue) Implements System.Collections.Generic.IDictionary(Of TKey, TValue).Values
            Get
                SyncLock Me.m_lock
                    Return Me._internalDictionary.Values
                End SyncLock
            End Get
        End Property

#End Region

#Region "IEnumerable"

        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 enumerator As New List(Of KeyValuePair(Of TKey, TValue))()

            ' Get keys
            Dim keyBuckets As Dictionary(Of TKey, TValue).KeyCollection = Me._internalDictionary.Keys
            If keyBuckets.Count > 0 Then
                ' Force memory access order
                Thread.MemoryBarrier()
                ' Fill list
                For i As Integer = 0 To keyBuckets.Count - 1
                    Dim value = Me._internalDictionary.Values(i)
                    enumerator.Add(New KeyValuePair(Of TKey, TValue)(keyBuckets(i), value))
                Next
            End If

            Return enumerator
        End Function

        Public Function GetEnumerator1() As System.Collections.IEnumerator Implements System.Collections.IEnumerable.GetEnumerator
            Return New DictionaryEnumerator(Me)
        End Function

        ''' <summary>
        ''' Private dictionary enumerator.
        ''' </summary>
        ''' <remarks></remarks>
        Private Class DictionaryEnumerator
            Implements IEnumerator, IDictionaryEnumerator

            Private m_enumerator As IEnumerator(Of KeyValuePair(Of TKey, TValue))

            Friend Sub New(ByVal dictionary As ThreadSafeDictionary(Of TKey, TValue))
                Me.m_enumerator = dictionary.GetEnumerator()
            End Sub

            Public ReadOnly Property Entry() As System.Collections.DictionaryEntry Implements System.Collections.IDictionaryEnumerator.Entry
                Get
                    Dim cKey = Me.m_enumerator.Current.Key
                    Dim cValue = Me.m_enumerator.Current.Value
                    Return New DictionaryEntry(cKey, cValue)
                End Get
            End Property

            Public ReadOnly Property Key() As Object Implements System.Collections.IDictionaryEnumerator.Key
                Get
                    Return Me.m_enumerator.Current.Key
                End Get
            End Property

            Public ReadOnly Property Value() As Object Implements System.Collections.IDictionaryEnumerator.Value
                Get
                    Return Me.m_enumerator.Current.Value
                End Get
            End Property

            Public ReadOnly Property Current() As Object Implements System.Collections.IEnumerator.Current
                Get
                    Return Me.Entry
                End Get
            End Property

            Public Function MoveNext() As Boolean Implements System.Collections.IEnumerator.MoveNext
                Return Me.m_enumerator.MoveNext()
            End Function

            Public Sub Reset() Implements System.Collections.IEnumerator.Reset
                Me.m_enumerator.Reset()
            End Sub
        End Class
#End Region

        ''' <summary>
        ''' Attempts to add the specified key and value to the dictionary.
        ''' </summary>
        ''' <param name="key">The key of the element to add.</param>
        ''' <param name="value">The value of the element to add.  The value can be null for reference types.</param>
        ''' <returns>True if the key/value pair was added to the dictionary successfully; otherwise, false if the key already exists.</returns>
        ''' <remarks></remarks>
        Public Function TryAdd(ByVal key As TKey, ByVal value As TValue) As Boolean
            If Me.ContainsKey(key) Then
                SyncLock Me.m_lock
                    If Me.ContainsKey(key) Then
                        Me._internalDictionary.Add(key, value)
                        Return True
                    End If
                End SyncLock
            End If
            Return False
        End Function

        ''' <summary>
        ''' Attempts to remove and return the value that has the specified key from the dictionary.
        ''' </summary>
        ''' <param name="key">The key of the element to remove and return.</param>
        ''' <param name="value">Out. Contains the object removed from the dictionary, 
        ''' or the default value of the TValue type if key does not exist.</param>
        ''' <returns>True if the object was removed successfully; otherwise, false.</returns>
        ''' <remarks></remarks>
        Public Function TryRemove(ByVal key As TKey, ByRef value As TValue) As Boolean
            If Me.ContainsKey(key) Then
                SyncLock Me.m_lock
                    If Me.ContainsKey(key) Then
                        ' Get value and remove
                        Dim current = Me._internalDictionary(key)
                        Dim flag = Me._internalDictionary.Remove(key)
                        ' Coerce value if true...
                        If (flag) Then _
                            value = current
                        Return flag
                    End If
                End SyncLock
            End If
            Return False
        End Function

        ''' <summary>
        ''' Compares the existing value for the specified key with a specified 
        ''' value, and if they are equal, updates the key with a third value.
        ''' </summary>
        ''' <param name="key">The key whose value is compared with comparisonValue and possibly replaced.</param>
        ''' <param name="newValue">The value that replaces the value of the element 
        ''' that has the specified key if the comparison results in equality.</param>
        ''' <param name="comparisonValue">The value that is compared to the 
        ''' value of the element that has the specified key.</param>
        ''' <returns>true if the value with key was equal to comparisonValue 
        ''' and was replaced with newValue; otherwise, false.</returns>
        ''' <remarks></remarks>
        Public Function TryUpdate(ByVal key As TKey, ByVal newValue As TValue, ByVal comparisonValue As TValue) As Boolean
            If (key Is Nothing) Then _
                Throw New ArgumentNullException("key")
            
            If Me.ContainsKey(key) Then
                Dim keyComp As IEqualityComparer(Of TKey) = Me._internalDictionary.Comparer
                Dim equalDef As IEqualityComparer(Of TValue) = EqualityComparer(Of TValue).Default
                Dim enumerator = Me.GetEnumerator()
                enumerator.Reset()
                Do While (enumerator.MoveNext())
                    Dim current = enumerator.Current
                    SyncLock Me.m_lock
                        If keyComp.Equals(current.Key, key) Then
                            If equalDef.Equals(current.Value, comparisonValue) Then
                                Me._internalDictionary.Item(key) = newValue
                                Return True
                            End If
                        End If
                    End SyncLock
                Loop
            End If

            Return False
        End Function

    End Class

End Namespace