17 files changed, 133 insertions, 32 deletions

diff --git a/core/tests/resourceflaggingtests/src/com/android/resourceflaggingtests/ResourceFlaggingTest.java b/core/tests/resourceflaggingtests/src/com/android/resourceflaggingtests/ResourceFlaggingTest.java
index ee701f674715..d9e90fa001c5 100644
--- a/core/tests/resourceflaggingtests/src/com/android/resourceflaggingtests/ResourceFlaggingTest.java
+++ b/core/tests/resourceflaggingtests/src/com/android/resourceflaggingtests/ResourceFlaggingTest.java
@@ -29,6 +29,7 @@ import android.util.DisplayMetrics;
 import android.view.LayoutInflater;
 import android.view.View;
 import android.widget.LinearLayout;
+import android.widget.TextView;
 
 import androidx.test.InstrumentationRegistry;
 import androidx.test.filters.SmallTest;
@@ -132,6 +133,24 @@ public class ResourceFlaggingTest {
         assertThat((View) ll.findViewById(R.id.text2)).isNotNull();
     }
 
+    @Test
+    public void testEnabledFlagLayoutOverrides() {
+        LinearLayout ll = (LinearLayout) getLayoutInflater().inflate(R.layout.layout3, null);
+        assertThat(ll).isNotNull();
+        assertThat((View) ll.findViewById(R.id.text1)).isNotNull();
+        assertThat(((TextView) ll.findViewById(R.id.text1)).getText()).isEqualTo("foobar");
+    }
+
+    @Test(expected = Resources.NotFoundException.class)
+    public void testDisabledLayout() {
+        getLayoutInflater().inflate(R.layout.layout2, null);
+    }
+
+    @Test(expected = Resources.NotFoundException.class)
+    public void testDisabledDrawable() {
+        mResources.getDrawable(R.drawable.removedpng);
+    }
+
     private LayoutInflater getLayoutInflater() {
         ContextWrapper c = new ContextWrapper(mContext) {
             private LayoutInflater mInflater;
diff --git a/tools/aapt2/Resource.h b/tools/aapt2/Resource.h
index 446fdd4dbc1b..0d261abd728d 100644
--- a/tools/aapt2/Resource.h
+++ b/tools/aapt2/Resource.h
@@ -243,6 +243,12 @@ struct ResourceFile {
 
   // Exported symbols
   std::vector<SourcedResourceName> exported_symbols;
+
+  // Flag status
+  FlagStatus flag_status = FlagStatus::NoFlag;
+
+  // Flag
+  std::optional<FeatureFlagAttribute> flag;
 };
 
 /**
diff --git a/tools/aapt2/ResourceParser.cpp b/tools/aapt2/ResourceParser.cpp
index a2383ac0285f..fce6aa7c80d9 100644
--- a/tools/aapt2/ResourceParser.cpp
+++ b/tools/aapt2/ResourceParser.cpp
@@ -547,7 +547,7 @@ bool ResourceParser::ParseResource(xml::XmlPullParser* parser,
   });
 
   std::string_view resource_type = parser->element_name();
-  if (auto flag = GetFlag(parser)) {
+  if (auto flag = ParseFlag(xml::FindAttribute(parser, xml::kSchemaAndroid, "featureFlag"))) {
     if (options_.flag) {
       diag_->Error(android::DiagMessage(source_.WithLine(parser->line_number()))
                    << "Resource flag are not allowed both in the path and in the file");
@@ -747,22 +747,6 @@ bool ResourceParser::ParseResource(xml::XmlPullParser* parser,
   return false;
 }
 
-std::optional<FeatureFlagAttribute> ResourceParser::GetFlag(xml::XmlPullParser* parser) {
-  auto name = xml::FindAttribute(parser, xml::kSchemaAndroid, "featureFlag");
-  if (name) {
-    FeatureFlagAttribute flag;
-    if (name->starts_with('!')) {
-      flag.negated = true;
-      flag.name = name->substr(1);
-    } else {
-      flag.name = name.value();
-    }
-    return flag;
-  } else {
-    return {};
-  }
-}
-
 bool ResourceParser::ParseItem(xml::XmlPullParser* parser,
                                ParsedResource* out_resource,
                                const uint32_t format) {
@@ -1669,7 +1653,7 @@ bool ResourceParser::ParseArrayImpl(xml::XmlPullParser* parser,
     const std::string& element_namespace = parser->element_namespace();
     const std::string& element_name = parser->element_name();
     if (element_namespace.empty() && element_name == "item") {
-      auto flag = GetFlag(parser);
+      auto flag = ParseFlag(xml::FindAttribute(parser, xml::kSchemaAndroid, "featureFlag"));
       std::unique_ptr<Item> item = ParseXml(parser, typeMask, kNoRawString);
       if (!item) {
         diag_->Error(android::DiagMessage(item_source) << "could not parse array item");
diff --git a/tools/aapt2/ResourceParser.h b/tools/aapt2/ResourceParser.h
index 4ad334c94518..90690d522ef2 100644
--- a/tools/aapt2/ResourceParser.h
+++ b/tools/aapt2/ResourceParser.h
@@ -90,8 +90,6 @@ class ResourceParser {
  private:
   DISALLOW_COPY_AND_ASSIGN(ResourceParser);
 
-  std::optional<FeatureFlagAttribute> GetFlag(xml::XmlPullParser* parser);
-
   std::optional<FlattenedXmlSubTree> CreateFlattenSubTree(xml::XmlPullParser* parser);
 
   // Parses the XML subtree as a StyleString (flattened XML representation for strings with
diff --git a/tools/aapt2/ResourcesInternal.proto b/tools/aapt2/ResourcesInternal.proto
index b0ed3da33368..f4735a2f6ce7 100644
--- a/tools/aapt2/ResourcesInternal.proto
+++ b/tools/aapt2/ResourcesInternal.proto
@@ -49,4 +49,9 @@ message CompiledFile {
 
   // Any symbols this file auto-generates/exports (eg. @+id/foo in an XML file).
   repeated Symbol exported_symbol = 5;
+
+  // The status of the flag the file is behind if any
+  uint32 flag_status = 6;
+  bool flag_negated = 7;
+  string flag_name = 8;
 }
diff --git a/tools/aapt2/cmd/Compile.cpp b/tools/aapt2/cmd/Compile.cpp
index cd8c187c1e99..52372fa38525 100644
--- a/tools/aapt2/cmd/Compile.cpp
+++ b/tools/aapt2/cmd/Compile.cpp
@@ -196,20 +196,11 @@ static bool CompileTable(IAaptContext* context, const CompileOptions& options,
     // If visibility was forced, we need to use it when creating a new resource and also error if
     // we try to parse the <public>, <public-group>, <java-symbol> or <symbol> tags.
     parser_options.visibility = options.visibility;
+    parser_options.flag = ParseFlag(path_data.flag_name);
 
-    if (!path_data.flag_name.empty()) {
-      FeatureFlagAttribute flag;
-      const auto& name = path_data.flag_name;
-      if (name.starts_with('!')) {
-        flag.negated = true;
-        flag.name = name.substr(1);
-      } else {
-        flag.name = name;
-      }
-      parser_options.flag = flag;
-
+    if (parser_options.flag) {
       std::string error;
-      auto flag_status = GetFlagStatus(flag, options.feature_flag_values, &error);
+      auto flag_status = GetFlagStatus(parser_options.flag, options.feature_flag_values, &error);
       if (flag_status) {
         parser_options.flag_status = std::move(flag_status.value());
       } else {
@@ -443,6 +434,18 @@ static bool CompileXml(IAaptContext* context, const CompileOptions& options,
   xmlres->file.config = path_data.config;
   xmlres->file.source = path_data.source;
   xmlres->file.type = ResourceFile::Type::kProtoXml;
+  xmlres->file.flag = ParseFlag(path_data.flag_name);
+
+  if (xmlres->file.flag) {
+    std::string error;
+    auto flag_status = GetFlagStatus(xmlres->file.flag, options.feature_flag_values, &error);
+    if (flag_status) {
+      xmlres->file.flag_status = flag_status.value();
+    } else {
+      context->GetDiagnostics()->Error(android::DiagMessage(path_data.source) << error);
+      return false;
+    }
+  }
 
   // Collect IDs that are defined here.
   XmlIdCollector collector;
@@ -532,6 +535,27 @@ static bool CompilePng(IAaptContext* context, const CompileOptions& options,
   res_file.source = path_data.source;
   res_file.type = ResourceFile::Type::kPng;
 
+  if (!path_data.flag_name.empty()) {
+    FeatureFlagAttribute flag;
+    auto name = path_data.flag_name;
+    if (name.starts_with('!')) {
+      flag.negated = true;
+      flag.name = name.substr(1);
+    } else {
+      flag.name = name;
+    }
+    res_file.flag = flag;
+
+    std::string error;
+    auto flag_status = GetFlagStatus(flag, options.feature_flag_values, &error);
+    if (flag_status) {
+      res_file.flag_status = flag_status.value();
+    } else {
+      context->GetDiagnostics()->Error(android::DiagMessage(path_data.source) << error);
+      return false;
+    }
+  }
+
   {
     auto data = file->OpenAsData();
     if (!data) {
diff --git a/tools/aapt2/cmd/Util.cpp b/tools/aapt2/cmd/Util.cpp
index 2177c345260c..08f8f0d85807 100644
--- a/tools/aapt2/cmd/Util.cpp
+++ b/tools/aapt2/cmd/Util.cpp
@@ -34,6 +34,20 @@ using ::android::base::StringPrintf;
 
 namespace aapt {
 
+std::optional<FeatureFlagAttribute> ParseFlag(std::optional<std::string_view> flag_text) {
+  if (!flag_text || flag_text->empty()) {
+    return {};
+  }
+  FeatureFlagAttribute flag;
+  if (flag_text->starts_with('!')) {
+    flag.negated = true;
+    flag.name = flag_text->substr(1);
+  } else {
+    flag.name = flag_text.value();
+  }
+  return flag;
+}
+
 std::optional<FlagStatus> GetFlagStatus(const std::optional<FeatureFlagAttribute>& flag,
                                         const FeatureFlagValues& feature_flag_values,
                                         std::string* out_err) {
diff --git a/tools/aapt2/cmd/Util.h b/tools/aapt2/cmd/Util.h
index f8e44b72247c..d32e532b86a8 100644
--- a/tools/aapt2/cmd/Util.h
+++ b/tools/aapt2/cmd/Util.h
@@ -49,6 +49,8 @@ struct FeatureFlagProperties {
 
 using FeatureFlagValues = std::map<std::string, FeatureFlagProperties, std::less<>>;
 
+std::optional<FeatureFlagAttribute> ParseFlag(std::optional<std::string_view> flag_text);
+
 std::optional<FlagStatus> GetFlagStatus(const std::optional<FeatureFlagAttribute>& flag,
                                         const FeatureFlagValues& feature_flag_values,
                                         std::string* out_err);
diff --git a/tools/aapt2/format/proto/ProtoDeserialize.cpp b/tools/aapt2/format/proto/ProtoDeserialize.cpp
index c6bd6dd3c4b2..8583cadff6d2 100644
--- a/tools/aapt2/format/proto/ProtoDeserialize.cpp
+++ b/tools/aapt2/format/proto/ProtoDeserialize.cpp
@@ -643,6 +643,12 @@ bool DeserializeCompiledFileFromPb(const pb::internal::CompiledFile& pb_file,
   out_file->source.path = pb_file.source_path();
   out_file->type = DeserializeFileReferenceTypeFromPb(pb_file.type());
 
+  out_file->flag_status = (FlagStatus)pb_file.flag_status();
+  if (!pb_file.flag_name().empty()) {
+    out_file->flag =
+        FeatureFlagAttribute{.name = pb_file.flag_name(), .negated = pb_file.flag_negated()};
+  }
+
   std::string config_error;
   if (!DeserializeConfigFromPb(pb_file.config(), &out_file->config, &config_error)) {
     std::ostringstream error;
diff --git a/tools/aapt2/format/proto/ProtoSerialize.cpp b/tools/aapt2/format/proto/ProtoSerialize.cpp
index 9c28780e7880..d83fe916ee95 100644
--- a/tools/aapt2/format/proto/ProtoSerialize.cpp
+++ b/tools/aapt2/format/proto/ProtoSerialize.cpp
@@ -755,6 +755,11 @@ void SerializeCompiledFileToPb(const ResourceFile& file, pb::internal::CompiledF
   out_file->set_source_path(file.source.path);
   out_file->set_type(SerializeFileReferenceTypeToPb(file.type));
   SerializeConfig(file.config, out_file->mutable_config());
+  out_file->set_flag_status((uint32_t)file.flag_status);
+  if (file.flag) {
+    out_file->set_flag_negated(file.flag->negated);
+    out_file->set_flag_name(file.flag->name);
+  }
 
   for (const SourcedResourceName& exported : file.exported_symbols) {
     pb::internal::CompiledFile_Symbol* pb_symbol = out_file->add_exported_symbol();
diff --git a/tools/aapt2/integration-tests/FlaggedResourcesTest/Android.bp b/tools/aapt2/integration-tests/FlaggedResourcesTest/Android.bp
index 432b84e5ff90..7160b35033da 100644
--- a/tools/aapt2/integration-tests/FlaggedResourcesTest/Android.bp
+++ b/tools/aapt2/integration-tests/FlaggedResourcesTest/Android.bp
@@ -31,12 +31,17 @@ genrule {
         "res/values/ints.xml",
         "res/values/strings.xml",
         "res/layout/layout1.xml",
+        "res/layout/layout3.xml",
         "res/flag(test.package.falseFlag)/values/bools.xml",
+        "res/flag(test.package.falseFlag)/layout/layout2.xml",
+        "res/flag(test.package.falseFlag)/drawable/removedpng.png",
+        "res/flag(test.package.trueFlag)/layout/layout3.xml",
         "res/values/flag(test.package.trueFlag)/bools.xml",
         "res/values/flag(!test.package.trueFlag)/bools.xml",
         "res/values/flag(!test.package.falseFlag)/bools.xml",
     ],
     out: [
+        "drawable_removedpng.(test.package.falseFlag).png.flat",
         "values_bools.arsc.flat",
         "values_bools.(test.package.falseFlag).arsc.flat",
         "values_bools.(test.package.trueFlag).arsc.flat",
@@ -46,6 +51,9 @@ genrule {
         "values_ints.arsc.flat",
         "values_strings.arsc.flat",
         "layout_layout1.xml.flat",
+        "layout_layout2.(test.package.falseFlag).xml.flat",
+        "layout_layout3.xml.flat",
+        "layout_layout3.(test.package.trueFlag).xml.flat",
     ],
     cmd: "$(location aapt2) compile $(in) -o $(genDir) " +
         "--feature-flags test.package.falseFlag:ro=false,test.package.trueFlag:ro=true",
diff --git a/tools/aapt2/integration-tests/FlaggedResourcesTest/res/flag(test.package.falseFlag)/drawable/removedpng.png b/tools/aapt2/integration-tests/FlaggedResourcesTest/res/flag(test.package.falseFlag)/drawable/removedpng.png
new file mode 100644
index 000000000000..8a9e6984be96
--- /dev/null
+++ b/tools/aapt2/integration-tests/FlaggedResourcesTest/res/flag(test.package.falseFlag)/drawable/removedpng.png
diff --git a/tools/aapt2/integration-tests/FlaggedResourcesTest/res/flag(test.package.falseFlag)/layout/layout2.xml b/tools/aapt2/integration-tests/FlaggedResourcesTest/res/flag(test.package.falseFlag)/layout/layout2.xml
new file mode 100644
index 000000000000..dec5de72925a
--- /dev/null
+++ b/tools/aapt2/integration-tests/FlaggedResourcesTest/res/flag(test.package.falseFlag)/layout/layout2.xml
@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="utf-8"?>
+<LinearLayout xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="match_parent"
+    android:layout_height="match_parent"
+    android:orientation="vertical" >
+</LinearLayout>
\ No newline at end of file
diff --git a/tools/aapt2/integration-tests/FlaggedResourcesTest/res/flag(test.package.trueFlag)/layout/layout3.xml b/tools/aapt2/integration-tests/FlaggedResourcesTest/res/flag(test.package.trueFlag)/layout/layout3.xml
new file mode 100644
index 000000000000..5aeee0ee1e28
--- /dev/null
+++ b/tools/aapt2/integration-tests/FlaggedResourcesTest/res/flag(test.package.trueFlag)/layout/layout3.xml
@@ -0,0 +1,10 @@
+<?xml version="1.0" encoding="utf-8"?>
+<LinearLayout xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="match_parent"
+    android:layout_height="match_parent"
+    android:orientation="vertical" >
+    <TextView android:id="@+id/text1"
+        android:layout_width="wrap_content"
+        android:layout_height="wrap_content"
+        android:text="foobar" />
+</LinearLayout>
\ No newline at end of file
diff --git a/tools/aapt2/integration-tests/FlaggedResourcesTest/res/layout/layout3.xml b/tools/aapt2/integration-tests/FlaggedResourcesTest/res/layout/layout3.xml
new file mode 100644
index 000000000000..dec5de72925a
--- /dev/null
+++ b/tools/aapt2/integration-tests/FlaggedResourcesTest/res/layout/layout3.xml
@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="utf-8"?>
+<LinearLayout xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="match_parent"
+    android:layout_height="match_parent"
+    android:orientation="vertical" >
+</LinearLayout>
\ No newline at end of file
diff --git a/tools/aapt2/link/FlaggedResources_test.cpp b/tools/aapt2/link/FlaggedResources_test.cpp
index 67d0c41f6b6f..629300838bbe 100644
--- a/tools/aapt2/link/FlaggedResources_test.cpp
+++ b/tools/aapt2/link/FlaggedResources_test.cpp
@@ -76,6 +76,10 @@ TEST_F(FlaggedResourcesTest, DisabledResourcesRemovedFromTable) {
 
   std::string output;
   DumpResourceTableToString(loaded_apk.get(), &output);
+  ASSERT_EQ(output.find("bool4"), std::string::npos);
+  ASSERT_EQ(output.find("str1"), std::string::npos);
+  ASSERT_EQ(output.find("layout2"), std::string::npos);
+  ASSERT_EQ(output.find("removedpng"), std::string::npos);
 }
 
 TEST_F(FlaggedResourcesTest, DisabledResourcesRemovedFromTableChunks) {
@@ -87,6 +91,8 @@ TEST_F(FlaggedResourcesTest, DisabledResourcesRemovedFromTableChunks) {
 
   ASSERT_EQ(output.find("bool4"), std::string::npos);
   ASSERT_EQ(output.find("str1"), std::string::npos);
+  ASSERT_EQ(output.find("layout2"), std::string::npos);
+  ASSERT_EQ(output.find("removedpng"), std::string::npos);
 }
 
 TEST_F(FlaggedResourcesTest, DisabledResourcesInRJava) {
diff --git a/tools/aapt2/link/TableMerger.cpp b/tools/aapt2/link/TableMerger.cpp
index d21697938c62..1bef5f8b17f6 100644
--- a/tools/aapt2/link/TableMerger.cpp
+++ b/tools/aapt2/link/TableMerger.cpp
@@ -377,6 +377,8 @@ bool TableMerger::MergeFile(const ResourceFile& file_desc, bool overlay, io::IFi
   file_ref->SetSource(file_desc.source);
   file_ref->type = file_desc.type;
   file_ref->file = file;
+  file_ref->SetFlagStatus(file_desc.flag_status);
+  file_ref->SetFlag(file_desc.flag);
 
   ResourceTablePackage* pkg = table.FindOrCreatePackage(file_desc.name.package);
   pkg->FindOrCreateType(file_desc.name.type)