1
+ package net .tascalate .concurrent .core ;
2
+
3
+ import java .lang .ref .Reference ;
4
+ import java .lang .ref .ReferenceQueue ;
5
+ import java .util .concurrent .ConcurrentHashMap ;
6
+ import java .util .concurrent .ConcurrentMap ;
7
+ import java .util .function .Function ;
8
+
9
+ class FunctionMemoization <K , V > implements Function <K , V > {
10
+ private final ConcurrentMap <K , Object > producerMutexes = new ConcurrentHashMap <>();
11
+ private final ConcurrentMap <Object , Object > valueMap = new ConcurrentHashMap <>();
12
+
13
+ private final Function <? super K , ? extends V > fn ;
14
+ private final ReferenceType keyRefType ;
15
+ private final ReferenceType valueRefType ;
16
+ private final ReferenceQueue <K > queue ;
17
+
18
+ FunctionMemoization (Function <? super K , ? extends V > fn ) {
19
+ this (ReferenceType .WEAK , ReferenceType .SOFT , fn );
20
+ }
21
+
22
+ FunctionMemoization (ReferenceType keyRefType , ReferenceType valueRefType , Function <? super K , ? extends V > fn ) {
23
+ this .fn = fn ;
24
+ this .keyRefType = keyRefType ;
25
+ this .valueRefType = valueRefType ;
26
+ this .queue = keyRefType .createKeyReferenceQueue ();
27
+ }
28
+
29
+ @ Override
30
+ public V apply (K key ) {
31
+ expungeStaleEntries ();
32
+
33
+ Object lookupKeyRef = keyRefType .createLookupKey (key );
34
+ Object valueRef ;
35
+
36
+ // Try to get a cached value.
37
+ valueRef = valueMap .get (lookupKeyRef );
38
+ V value ;
39
+
40
+ if (valueRef != null ) {
41
+ value = valueRefType .dereference (valueRef );
42
+ if (value != null ) {
43
+ // A cached value was found.
44
+ return value ;
45
+ }
46
+ }
47
+
48
+ Object mutex = getOrCreateMutex (key );
49
+ synchronized (mutex ) {
50
+ try {
51
+ // Double-check after getting mutex
52
+ valueRef = valueMap .get (lookupKeyRef );
53
+ value = valueRef == null ? null : valueRefType .dereference (valueRef );
54
+ if (value == null ) {
55
+ value = fn .apply (key );
56
+ valueMap .put (
57
+ keyRefType .createKeyReference (key , queue ),
58
+ valueRefType .createValueReference (value )
59
+ );
60
+ }
61
+ } finally {
62
+ producerMutexes .remove (key , mutex );
63
+ }
64
+ }
65
+
66
+ return value ;
67
+ }
68
+
69
+ public V forget (K key ) {
70
+ Object mutex = getOrCreateMutex (key );
71
+ synchronized (mutex ) {
72
+ try {
73
+ Object valueRef = valueMap .remove (keyRefType .createLookupKey (key ));
74
+ return valueRef == null ? null : valueRefType .dereference (valueRef );
75
+ } finally {
76
+ producerMutexes .remove (key , mutex );
77
+ }
78
+ }
79
+ }
80
+
81
+ private Object getOrCreateMutex (K key ) {
82
+ Object createdMutex = new byte [0 ];
83
+ Object existingMutex = producerMutexes .putIfAbsent (key , createdMutex );
84
+ if (existingMutex != null ) {
85
+ return existingMutex ;
86
+ } else {
87
+ return createdMutex ;
88
+ }
89
+ }
90
+
91
+ private void expungeStaleEntries () {
92
+ for (Reference <? extends K > ref ; (ref = queue .poll ()) != null ;) {
93
+ @ SuppressWarnings ("unchecked" )
94
+ Reference <K > keyRef = (Reference <K >) ref ;
95
+ // keyRef now is equal only to itself while referent is cleared already
96
+ // so it's safe to remove it without ceremony (like getOrCreateMutex(keyRef) usage)
97
+ valueMap .remove (keyRef );
98
+ }
99
+ }
100
+ }
0 commit comments