From a9165797d9a6d08829bd290dc6e264d2fb448ad0 Mon Sep 17 00:00:00 2001
From: Isira Seneviratne <isirasen96@gmail.com>
Date: Sat, 26 Nov 2022 16:43:39 +0530
Subject: [PATCH] Use ConcurrentSkipListSet in PrettyTime.

---
 .../org/ocpsoft/prettytime/PrettyTime.java    | 51 +++++++++----------
 1 file changed, 23 insertions(+), 28 deletions(-)

diff --git a/core/src/main/java/org/ocpsoft/prettytime/PrettyTime.java b/core/src/main/java/org/ocpsoft/prettytime/PrettyTime.java
index 93b97eb..39385c7 100644
--- a/core/src/main/java/org/ocpsoft/prettytime/PrettyTime.java
+++ b/core/src/main/java/org/ocpsoft/prettytime/PrettyTime.java
@@ -30,7 +30,9 @@
 import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentSkipListSet;
 
 import org.ocpsoft.prettytime.impl.DurationImpl;
 import org.ocpsoft.prettytime.impl.ResourcesTimeFormat;
@@ -72,9 +74,10 @@ public class PrettyTime
 {
    private volatile Instant reference;
    private volatile Locale locale = Locale.getDefault();
-   private final Map<TimeUnit, TimeFormat> units = new ConcurrentHashMap<>();
-   private volatile List<TimeUnit> cachedUnits;
-   private String overrideResourceBundle;
+   private final Map<TimeUnit, TimeFormat> unitsAndFormats = new ConcurrentHashMap<>();
+   private final Set<TimeUnit> cachedUnits
+           = new ConcurrentSkipListSet<>(Comparator.comparingLong(TimeUnit::getMillisPerUnit));
+   private final String overrideResourceBundle;
 
    /**
     * Create a new {@link PrettyTime} instance that will always use the current value of
@@ -1306,12 +1309,12 @@ public String formatDurationUnrounded(final LocalDate then)
    public TimeFormat getFormat(TimeUnit unit)
    {
       if (unit == null) return null;
-      if (units.get(unit) != null) {
-         return units.get(unit);
+      if (unitsAndFormats.get(unit) != null) {
+         return unitsAndFormats.get(unit);
       } else {
          // Trying to transform the TimeUnit to String does the trick
          Map<String, TimeFormat> map = new ConcurrentHashMap<>();
-         units.keySet().forEach(key -> map.put(key.toString(), units.get(key)));
+         unitsAndFormats.keySet().forEach(key -> map.put(key.toString(), unitsAndFormats.get(key)));
          return map.get(unit.toString());
       }
    }
@@ -1438,13 +1441,7 @@ public PrettyTime setReference(final LocalDate localDate, final ZoneId zoneId)
     */
    public List<TimeUnit> getUnits()
    {
-      if (cachedUnits == null) {
-         List<TimeUnit> result = new ArrayList<>(units.keySet());
-         Collections.sort(result, Comparator.comparing(TimeUnit::getMillisPerUnit));
-         cachedUnits = Collections.unmodifiableList(result);
-      }
-
-      return cachedUnits;
+      return Collections.unmodifiableList(new ArrayList<>(cachedUnits));
    }
 
    /**
@@ -1456,7 +1453,7 @@ public <UNIT extends TimeUnit> UNIT getUnit(final Class<UNIT> unitType)
       if (unitType == null)
          return null;
 
-      for (TimeUnit unit : units.keySet()) {
+      for (TimeUnit unit : unitsAndFormats.keySet()) {
          if (unitType.isAssignableFrom(unit.getClass())) {
             return (UNIT) unit;
          }
@@ -1471,10 +1468,9 @@ public <UNIT extends TimeUnit> UNIT getUnit(final Class<UNIT> unitType)
     */
    public PrettyTime registerUnit(final TimeUnit unit, TimeFormat format)
    {
-      cachedUnits = null;
+      cachedUnits.add(unit);
+      unitsAndFormats.put(unit, format);
 
-      units.put(Objects.requireNonNull(unit, "TimeUnit to register must not be null."),
-               Objects.requireNonNull(format, "TimeFormat to register must not be null."));
       if (unit instanceof LocaleAware)
          ((LocaleAware<?>) unit).setLocale(locale);
       if (format instanceof LocaleAware)
@@ -1520,11 +1516,11 @@ public <UNIT extends TimeUnit> TimeFormat removeUnit(final Class<UNIT> unitType)
       if (unitType == null)
          return null;
 
-      for (TimeUnit unit : units.keySet()) {
+      for (TimeUnit unit : unitsAndFormats.keySet()) {
          if (unitType.isAssignableFrom(unit.getClass())) {
-            cachedUnits = null;
+            cachedUnits.remove(unit);
 
-            return units.remove(unit);
+            return unitsAndFormats.remove(unit);
          }
       }
       return null;
@@ -1540,9 +1536,9 @@ public TimeFormat removeUnit(final TimeUnit unit)
       if (unit == null)
          return null;
 
-      cachedUnits = null;
+      cachedUnits.remove(unit);
 
-      return units.remove(unit);
+      return unitsAndFormats.remove(unit);
    }
 
    /**
@@ -1563,15 +1559,14 @@ public PrettyTime setLocale(Locale locale)
          locale = Locale.getDefault();
 
       this.locale = locale;
-      for (TimeUnit unit : units.keySet()) {
+      for (TimeUnit unit : unitsAndFormats.keySet()) {
          if (unit instanceof LocaleAware)
             ((LocaleAware<?>) unit).setLocale(locale);
       }
-      for (TimeFormat format : units.values()) {
+      for (TimeFormat format : unitsAndFormats.values()) {
          if (format instanceof LocaleAware)
             ((LocaleAware<?>) format).setLocale(locale);
       }
-      cachedUnits = null;
       return this;
    }
 
@@ -1586,9 +1581,9 @@ public String toString()
     */
    public List<TimeUnit> clearUnits()
    {
-      List<TimeUnit> result = getUnits();
-      cachedUnits = null;
-      units.clear();
+      List<TimeUnit> result = new ArrayList<>(cachedUnits);
+      cachedUnits.clear();
+      unitsAndFormats.clear();
       return result;
    }