33import org .comroid .api .Junction ;
44import org .comroid .api .Polyfill ;
55import org .comroid .mutatio .ref .Reference ;
6+ import org .comroid .mutatio .ref .ReferenceIndex ;
7+ import org .comroid .mutatio .ref .ReferenceMap ;
68import org .jetbrains .annotations .NotNull ;
79import org .jetbrains .annotations .Nullable ;
810
1214import java .util .stream .Collectors ;
1315import java .util .stream .Stream ;
1416
15- public interface TrieMap <K , V > extends Map <K , V > {
17+ public interface TrieMap <K , V > extends ReferenceMap < K , V , Reference . Settable < V >>, Map <K , V > {
1618 Junction <K , String > getKeyConverter ();
1719
1820 @ Override
@@ -39,6 +41,11 @@ default void putAll(@NotNull Map<? extends K, ? extends V> map) {
3941
4042 Stage <V > getStage (String key );
4143
44+ default V get (Object key ) {
45+ //noinspection unchecked
46+ return getReference ((K ) key ).get ();
47+ }
48+
4249 final class Stage <V > implements Map .Entry <String , V > {
4350 private final Map <Character , Stage <V >> storage = new ConcurrentHashMap <>();
4451 private final Reference .Settable <V > reference = Reference .Settable .create ();
@@ -90,48 +97,53 @@ private Stream<Stage<V>> streamPresentStages() {
9097 );
9198 }
9299
93- private Optional <V > getValue (char [] chars , int cIndex ) {
100+ private Optional <Reference . Settable < V >> getReference (char [] chars , int cIndex ) {
94101 if (cIndex >= chars .length )
95- return reference . wrap ( );
102+ return Optional . of ( reference );
96103
97- return Optional . ofNullable ( storage . getOrDefault ( chars [cIndex ], null ) )
98- .flatMap (stage -> stage .getValue (chars , cIndex + 1 ));
104+ return getStageByChar ( chars [cIndex ])
105+ .flatMap (stage -> stage .getReference (chars , cIndex + 1 ));
99106 }
100107
101108 private Optional <V > putValue (char [] chars , int cIndex , @ Nullable V value ) {
102109 if (cIndex >= chars .length )
103110 return Optional .ofNullable (setValue (value ));
104111
105112 // expect existing stages
106- return Optional . ofNullable ( storage . getOrDefault ( chars [cIndex ], null ) )
113+ return getStageByChar ( chars [cIndex ])
107114 .flatMap (stage -> stage .putValue (chars , cIndex + 1 , value ));
108115 }
109116
110117 private Optional <V > remove (char [] chars , int cIndex ) {
111118 if (cIndex >= chars .length )
112119 return Optional .ofNullable (remove ());
113120
114- return Optional . ofNullable ( storage . getOrDefault ( chars [cIndex ], null ) )
121+ return getStageByChar ( chars [cIndex ])
115122 .flatMap (stage -> stage .remove (chars , cIndex ));
116123 }
117124
118125 public boolean containsKey (char [] chars , int cIndex ) {
119126 if (cIndex >= chars .length )
120127 return false ;
121128
122- return Optional . ofNullable ( storage . getOrDefault ( chars [cIndex ], null ) )
129+ return getStageByChar ( chars [cIndex ])
123130 .map (stage -> stage .containsKey (chars , cIndex + 1 ))
124131 .orElse (false );
125132 }
126133
127- private Stage <V > requireStage (char [] chars , int cIndex ) {
134+ public Stage <V > requireStage (char [] chars , int cIndex ) {
128135 if (cIndex < chars .length ) {
129136 return storage .computeIfAbsent (chars [cIndex ], key -> {
130137 String converted = new String (Arrays .copyOfRange (chars , 0 , cIndex + 1 ));
131138 return new Stage <>(converted );
132139 }).requireStage (chars , cIndex + 1 );
133140 } else return this ;
134141 }
142+
143+ @ NotNull
144+ public Optional <Stage <V >> getStageByChar (char aChar ) {
145+ return Optional .ofNullable (storage .getOrDefault (aChar , null ));
146+ }
135147 }
136148
137149 class Basic <K , V > implements TrieMap <K , V > {
@@ -175,13 +187,58 @@ public boolean containsValue(Object value) {
175187 }
176188
177189 @ Override
178- public V get ( Object key ) {
190+ public Reference . Settable < V > getReference ( K key , boolean createIfAbsent ) {
179191 final char [] convertKey = convertKey (key );
180192
181193 if (convertKey .length == 0 )
182- return null ;
194+ return Reference .Settable .create ();
195+
196+ return baseStage .getReference (convertKey , 0 )
197+ .orElseGet (Reference .Settable ::create );
198+ }
183199
184- return baseStage .getValue (convertKey , 0 ).orElse (null );
200+ @ Override
201+ public ReferenceIndex <Entry <K , V >> entryIndex () {
202+ class RemoteIndex implements ReferenceIndex <Entry <K , V >> {
203+ private final ArrayList <Entry <K , V >> entries = new ArrayList <>(entrySet ());
204+
205+ @ Override
206+ public List <Entry <K , V >> unwrap () {
207+ return entries ;
208+ }
209+
210+ @ Override
211+ public int size () {
212+ return entries .size ();
213+ }
214+
215+ @ Override
216+ public boolean add (Entry <K , V > entry ) {
217+ Basic .this .put (entry .getKey (), entry .getValue ());
218+ return Basic .this .containsKey (entry .getKey ());
219+ }
220+
221+ @ Override
222+ public boolean remove (Entry <K , V > entry ) {
223+ Basic .this .remove (entry .getKey ());
224+ return !Basic .this .containsKey (entry .getKey ());
225+ }
226+
227+ @ Override
228+ public void clear () {
229+ Basic .this .clear ();
230+ }
231+
232+ @ Override
233+ public Reference <Entry <K , V >> getReference (int index ) {
234+ return Reference .conditional (
235+ () -> entries .size () < index ,
236+ () -> entries .get (index )
237+ );
238+ }
239+ }
240+
241+ return new RemoteIndex ();
185242 }
186243
187244 @ Nullable
0 commit comments