A new function to look up time zone ID info am: d067fd1cda

Original change: https://android-review.googlesource.com/c/platform/system/timezone/+/1398076

Change-Id: Ife78c1c2bb19144bb6d9bb362ff052c726fae57c
diff --git a/input_tools/android/tzids/src/main/java/com/android/timezone/tzids/TimeZoneIds.java b/input_tools/android/tzids/src/main/java/com/android/timezone/tzids/TimeZoneIds.java
index bdf9e5e..cc35e17 100644
--- a/input_tools/android/tzids/src/main/java/com/android/timezone/tzids/TimeZoneIds.java
+++ b/input_tools/android/tzids/src/main/java/com/android/timezone/tzids/TimeZoneIds.java
@@ -104,6 +104,33 @@
     }
 
     /**
+     * Returns the country code associated with the supplied {@code zoneId} or {@code null} if it
+     * is not found.
+     */
+    public String getCountryCodeForZoneId(String zoneId) {
+        for (TzIdsProto.CountryMapping countryMapping
+                : mTimeZoneIdsProto.getCountryMappingsList()) {
+            if (countryMapping.getTimeZoneIdsList().contains(zoneId)) {
+                return countryMapping.getIsoCode();
+            }
+
+            for (TzIdsProto.TimeZoneLink timeZoneLink : countryMapping.getTimeZoneLinksList()) {
+                if (timeZoneLink.getAlternativeId().equals(zoneId)) {
+                    return countryMapping.getIsoCode();
+                }
+            }
+
+            for (TzIdsProto.TimeZoneReplacement timeZoneReplacement
+                    : countryMapping.getTimeZoneReplacementsList()) {
+                if (timeZoneReplacement.getReplacedId().equals(zoneId)) {
+                    return countryMapping.getIsoCode();
+                }
+            }
+        }
+        return null;
+    }
+
+    /**
      * Returns a country mapping for TZDB time zone IDs on or after the specified time.
      * The {@link Map} returned contains entries that link Olson IDs to themselves, or to preferred
      * Olson IDs. The {@code replacementThreshold} is used to identify which "replacements" should
diff --git a/input_tools/android/tzids/src/test/java/com/android/timezone/tzids/TimeZoneIdsTest.java b/input_tools/android/tzids/src/test/java/com/android/timezone/tzids/TimeZoneIdsTest.java
index 6dbeb2f..7d8fb63 100644
--- a/input_tools/android/tzids/src/test/java/com/android/timezone/tzids/TimeZoneIdsTest.java
+++ b/input_tools/android/tzids/src/test/java/com/android/timezone/tzids/TimeZoneIdsTest.java
@@ -18,6 +18,7 @@
 
 import static java.time.ZoneOffset.UTC;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
 
 import com.android.timezone.tzids.proto.TzIdsProto;
 import org.junit.Test;
@@ -114,6 +115,36 @@
         }
     }
 
+    @Test
+    public void getCountryCodeForZoneId() {
+        TzIdsProto.TimeZoneIds.Builder tzIdsBuilder = TzIdsProto.TimeZoneIds.newBuilder();
+
+        TzIdsProto.CountryMapping gb = TzIdsProto.CountryMapping.newBuilder()
+                .setIsoCode("gb")
+                .addTimeZoneIds("Europe/London")
+                .addTimeZoneLinks(createLink("GB", "Europe/London"))
+                .build();
+        tzIdsBuilder.addCountryMappings(gb);
+
+        Instant boiseFrom = LocalDateTime.of(1974, Month.FEBRUARY, 3, 9, 0).toInstant(UTC);
+        TzIdsProto.CountryMapping us = TzIdsProto.CountryMapping.newBuilder()
+                .setIsoCode("us")
+                .addTimeZoneIds("America/Phoenix")
+                .addTimeZoneLinks(createLink("US/Arizona", "America/Phoenix"))
+                .addTimeZoneReplacements(
+                        createReplacement("America/Boise", "America/Phoenix", boiseFrom))
+                .build();
+        tzIdsBuilder.addCountryMappings(us);
+
+        TimeZoneIds tzIds = new TimeZoneIds(tzIdsBuilder.build());
+        assertNull(tzIds.getCountryCodeForZoneId("FooBar"));
+        assertEquals("gb", tzIds.getCountryCodeForZoneId("GB"));
+        assertEquals("gb", tzIds.getCountryCodeForZoneId("Europe/London"));
+        assertEquals("us", tzIds.getCountryCodeForZoneId("America/Phoenix"));
+        assertEquals("us", tzIds.getCountryCodeForZoneId("US/Arizona"));
+        assertEquals("us", tzIds.getCountryCodeForZoneId("America/Boise"));
+    }
+
     private static TzIdsProto.TimeZoneLink createLink(
             String alternativeId, String preferredId) {
         return TzIdsProto.TimeZoneLink.newBuilder()