1 /*
2  * Copyright (C) 2019 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package com.android.server.utils.quota;
18 
19 import android.annotation.NonNull;
20 import android.annotation.Nullable;
21 import android.util.ArrayMap;
22 import android.util.SparseArrayMap;
23 
24 import java.util.function.Consumer;
25 import java.util.function.Function;
26 
27 /**
28  * A SparseArrayMap of ArrayMaps, which is suitable for holding userId-packageName-tag combination
29  * (UPTC)->object associations. Tags are any desired String.
30  *
31  * @see Uptc
32  */
33 class UptcMap<T> {
34     private final SparseArrayMap<String, ArrayMap<String, T>> mData = new SparseArrayMap<>();
35 
add(int userId, @NonNull String packageName, @Nullable String tag, @Nullable T obj)36     public void add(int userId, @NonNull String packageName, @Nullable String tag,
37             @Nullable T obj) {
38         ArrayMap<String, T> data = mData.get(userId, packageName);
39         if (data == null) {
40             data = new ArrayMap<>();
41             mData.add(userId, packageName, data);
42         }
43         data.put(tag, obj);
44     }
45 
clear()46     public void clear() {
47         mData.clear();
48     }
49 
contains(int userId, @NonNull String packageName)50     public boolean contains(int userId, @NonNull String packageName) {
51         return mData.contains(userId, packageName);
52     }
53 
contains(int userId, @NonNull String packageName, @Nullable String tag)54     public boolean contains(int userId, @NonNull String packageName, @Nullable String tag) {
55         // This structure never inserts a null ArrayMap, so if get(userId, packageName) returns
56         // null, the UPTC was never inserted.
57         ArrayMap<String, T> data = mData.get(userId, packageName);
58         return data != null && data.containsKey(tag);
59     }
60 
61     /** Removes all the data for the user, if there was any. */
delete(int userId)62     public void delete(int userId) {
63         mData.delete(userId);
64     }
65 
66     /** Removes the data for the user, package, and tag, if there was any. */
delete(int userId, @NonNull String packageName, @Nullable String tag)67     public void delete(int userId, @NonNull String packageName, @Nullable String tag) {
68         final ArrayMap<String, T> data = mData.get(userId, packageName);
69         if (data != null) {
70             data.remove(tag);
71             if (data.size() == 0) {
72                 mData.delete(userId, packageName);
73             }
74         }
75     }
76 
77     /** Removes the data for the user and package, if there was any. */
delete(int userId, @NonNull String packageName)78     public ArrayMap<String, T> delete(int userId, @NonNull String packageName) {
79         return mData.delete(userId, packageName);
80     }
81 
82     /**
83      * Returns the set of tag -> object mappings for the given userId and packageName
84      * combination.
85      */
86     @Nullable
get(int userId, @NonNull String packageName)87     public ArrayMap<String, T> get(int userId, @NonNull String packageName) {
88         return mData.get(userId, packageName);
89     }
90 
91     /** Returns the saved object for the given UPTC. */
92     @Nullable
get(int userId, @NonNull String packageName, @Nullable String tag)93     public T get(int userId, @NonNull String packageName, @Nullable String tag) {
94         final ArrayMap<String, T> data = mData.get(userId, packageName);
95         return data != null ? data.get(tag) : null;
96     }
97 
98     /**
99      * Returns the saved object for the given UPTC. If there was no saved object, it will create a
100      * new object using creator, insert it, and return it.
101      */
102     @Nullable
getOrCreate(int userId, @NonNull String packageName, @Nullable String tag, Function<Void, T> creator)103     public T getOrCreate(int userId, @NonNull String packageName, @Nullable String tag,
104             Function<Void, T> creator) {
105         final ArrayMap<String, T> data = mData.get(userId, packageName);
106         if (data == null || !data.containsKey(tag)) {
107             // We've never inserted data for this combination before. Create a new object.
108             final T val = creator.apply(null);
109             add(userId, packageName, tag, val);
110             return val;
111         }
112         return data.get(tag);
113     }
114 
115     /** Returns the userId at the given index. */
getUserIdAtIndex(int index)116     private int getUserIdAtIndex(int index) {
117         return mData.keyAt(index);
118     }
119 
120     /** Returns the package name at the given index. */
121     @NonNull
getPackageNameAtIndex(int userIndex, int packageIndex)122     private String getPackageNameAtIndex(int userIndex, int packageIndex) {
123         return mData.keyAt(userIndex, packageIndex);
124     }
125 
126     /** Returns the tag at the given index. */
127     @NonNull
getTagAtIndex(int userIndex, int packageIndex, int tagIndex)128     private String getTagAtIndex(int userIndex, int packageIndex, int tagIndex) {
129         // This structure never inserts a null ArrayMap, so if the indices are valid, valueAt()
130         // won't return null.
131         return mData.valueAt(userIndex, packageIndex).keyAt(tagIndex);
132     }
133 
134     /** Returns the size of the outer (userId) array. */
userCount()135     public int userCount() {
136         return mData.numMaps();
137     }
138 
139     /** Returns the number of packages saved for a given userId. */
packageCountForUser(int userId)140     public int packageCountForUser(int userId) {
141         return mData.numElementsForKey(userId);
142     }
143 
144     /** Returns the number of tags saved for a given userId-packageName combination. */
tagCountForUserAndPackage(int userId, @NonNull String packageName)145     public int tagCountForUserAndPackage(int userId, @NonNull String packageName) {
146         final ArrayMap data = mData.get(userId, packageName);
147         return data != null ? data.size() : 0;
148     }
149 
150     /** Returns the value T at the given user, package, and tag indices. */
151     @Nullable
valueAt(int userIndex, int packageIndex, int tagIndex)152     public T valueAt(int userIndex, int packageIndex, int tagIndex) {
153         final ArrayMap<String, T> data = mData.valueAt(userIndex, packageIndex);
154         return data != null ? data.valueAt(tagIndex) : null;
155     }
156 
forEach(Consumer<T> consumer)157     public void forEach(Consumer<T> consumer) {
158         mData.forEach((tagMap) -> {
159             for (int i = tagMap.size() - 1; i >= 0; --i) {
160                 consumer.accept(tagMap.valueAt(i));
161             }
162         });
163     }
164 
forEach(UptcDataConsumer<T> consumer)165     public void forEach(UptcDataConsumer<T> consumer) {
166         final int uCount = userCount();
167         for (int u = 0; u < uCount; ++u) {
168             final int userId = getUserIdAtIndex(u);
169 
170             final int pkgCount = packageCountForUser(userId);
171             for (int p = 0; p < pkgCount; ++p) {
172                 final String pkgName = getPackageNameAtIndex(u, p);
173 
174                 final int tagCount = tagCountForUserAndPackage(userId, pkgName);
175                 for (int t = 0; t < tagCount; ++t) {
176                     final String tag = getTagAtIndex(u, p, t);
177                     consumer.accept(userId, pkgName, tag, get(userId, pkgName, tag));
178                 }
179             }
180         }
181     }
182 
183     interface UptcDataConsumer<D> {
accept(int userId, @NonNull String packageName, @Nullable String tag, @Nullable D obj)184         void accept(int userId, @NonNull String packageName, @Nullable String tag, @Nullable D obj);
185     }
186 }
187