Create NotificationOptimizedLinearLayout

This CL introduces a custom LinearLayout class that overrides the onMeasure method.
For a specific layout pattern (a LinearLayout, whose width is match_parent, with one weighted child), this custom class will strategically measure unweighted views first
 and then measure the weighted view with remaining space.
This allows it to calculate the correct layout dimensions in a single pass, avoiding the performance overhead of the double measurement

Bug: 316110233
Test: presubmit + silver test + screenshot test of Messaging Notifications
Flag: ACONFIG notif_linearlayout_optimized DEVELOPMENT
Change-Id: I012824cb9cb24c95e5d29d8a3a9f8dde359fef96

6 files changed, 960 insertions, 3 deletions

diff --git a/core/java/com/android/internal/widget/MessagingGroup.java b/core/java/com/android/internal/widget/MessagingGroup.java
index 30e4099f9a6f..f8f104943c61 100644
--- a/core/java/com/android/internal/widget/MessagingGroup.java
+++ b/core/java/com/android/internal/widget/MessagingGroup.java
@@ -55,7 +55,9 @@ import java.util.List;
  * A message of a {@link MessagingLayout}.
  */
 @RemoteViews.RemoteView
-public class MessagingGroup extends LinearLayout implements MessagingLinearLayout.MessagingChild {
+public class MessagingGroup extends NotificationOptimizedLinearLayout implements
+        MessagingLinearLayout.MessagingChild {
+
     private static final MessagingPool<MessagingGroup> sInstancePool =
             new MessagingPool<>(10);
 
diff --git a/core/java/com/android/internal/widget/NotificationOptimizedLinearLayout.java b/core/java/com/android/internal/widget/NotificationOptimizedLinearLayout.java
new file mode 100644
index 000000000000..b5e9b8f537e4
--- /dev/null
+++ b/core/java/com/android/internal/widget/NotificationOptimizedLinearLayout.java
@@ -0,0 +1,578 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.internal.widget;
+
+import static android.widget.flags.Flags.notifLinearlayoutOptimized;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.content.Context;
+import android.content.res.Resources;
+import android.graphics.drawable.Drawable;
+import android.os.Build;
+import android.os.Trace;
+import android.util.AttributeSet;
+import android.util.Log;
+import android.view.Gravity;
+import android.view.View;
+import android.view.ViewGroup;
+import android.widget.LinearLayout;
+import android.widget.RemoteViews;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * This LinearLayout customizes the measurement behavior of LinearLayout for Notification layouts.
+ * When there is exactly
+ * one child View with <code>layout_weight</code>. onMeasure methods of this LinearLayout will:
+ * 1. Measure all other children.
+ * 2. Calculate the remaining space for the View with <code>layout_weight</code>
+ * 3. Measure the weighted View using the calculated remaining width or height (based on
+ * Orientation).
+ * This ensures that the weighted View fills the remaining space in LinearLayout with only single
+ * measure.
+ *
+ * **Assumptions:**
+ * - There is *exactly one* child view with non-zero <code>layout_weight</code>.
+ * - Other views should not have weight.
+ * - LinearLayout doesn't have <code>weightSum</code>.
+ * - Horizontal LinearLayout's width should be measured EXACTLY.
+ * - Horizontal LinearLayout shouldn't need baseLineAlignment.
+ * - Vertical LinearLayout shouldn't have MATCH_PARENT children when it is not measured EXACTLY.
+ *
+ * @hide
+ */
+@RemoteViews.RemoteView
+public class NotificationOptimizedLinearLayout extends LinearLayout {
+    private static final boolean DEBUG_LAYOUT = false;
+    private static final boolean TRACE_ONMEASURE = Build.isDebuggable();
+    private static final String TAG = "NotifOptimizedLinearLayout";
+
+    private boolean mShouldUseOptimizedLayout = false;
+
+    public NotificationOptimizedLinearLayout(Context context) {
+        super(context);
+    }
+
+    public NotificationOptimizedLinearLayout(Context context, @Nullable AttributeSet attrs) {
+        super(context, attrs);
+    }
+
+    public NotificationOptimizedLinearLayout(Context context, @Nullable AttributeSet attrs,
+            int defStyleAttr) {
+        super(context, attrs, defStyleAttr);
+    }
+
+    public NotificationOptimizedLinearLayout(Context context, AttributeSet attrs, int defStyleAttr,
+            int defStyleRes) {
+        super(context, attrs, defStyleAttr, defStyleRes);
+    }
+
+    @Override
+    protected void onMeasure(int widthMeasureSpec, int heightMeasureSpec) {
+        final View weightedChildView = getSingleWeightedChild();
+        mShouldUseOptimizedLayout =
+                isUseOptimizedLinearLayoutFlagEnabled() && weightedChildView != null
+                        && isLinearLayoutUsable(widthMeasureSpec, heightMeasureSpec);
+
+        if (mShouldUseOptimizedLayout) {
+            onMeasureOptimized(weightedChildView, widthMeasureSpec, heightMeasureSpec);
+        } else {
+            super.onMeasure(widthMeasureSpec, heightMeasureSpec);
+        }
+    }
+
+    private boolean isUseOptimizedLinearLayoutFlagEnabled() {
+        final boolean enabled = notifLinearlayoutOptimized();
+        if (!enabled) {
+            logSkipOptimizedOnMeasure("enableNotifLinearlayoutOptimized flag is off.");
+        }
+        return enabled;
+    }
+
+    /**
+     * Checks if optimizations can be safely applied to this LinearLayout during layout
+     * calculations. Optimizations might be disabled in the following cases:
+     *
+     * **weightSum**: When LinearLayout has weightSum
+     * ** MATCH_PARENT children in non EXACT dimension**
+     * **Horizontal LinearLayout with non-EXACT width**
+     * **Baseline Alignment:**  If views need to align their baselines in Horizontal LinearLayout
+     *
+     * @param widthMeasureSpec  The width measurement specification.
+     * @param heightMeasureSpec The height measurement specification.
+     * @return `true` if optimization is possible, `false` otherwise.
+     */
+    private boolean isLinearLayoutUsable(int widthMeasureSpec, int heightMeasureSpec) {
+        final boolean hasWeightSum = getWeightSum() > 0.0f;
+        if (hasWeightSum) {
+            logSkipOptimizedOnMeasure("Has weightSum.");
+            return false;
+        }
+
+        if (requiresMatchParentRemeasureForVerticalLinearLayout(widthMeasureSpec)) {
+            logSkipOptimizedOnMeasure(
+                    "Vertical LinearLayout requires children width MATCH_PARENT remeasure ");
+            return false;
+        }
+
+        final boolean isHorizontal = getOrientation() == HORIZONTAL;
+        if (isHorizontal && MeasureSpec.getMode(widthMeasureSpec) != MeasureSpec.EXACTLY) {
+            logSkipOptimizedOnMeasure("Horizontal LinearLayout's width should be "
+                    + "measured EXACTLY");
+            return false;
+        }
+
+        if (requiresBaselineAlignmentForHorizontalLinearLayout()) {
+            logSkipOptimizedOnMeasure("Need to apply baseline.");
+            return false;
+        }
+        return true;
+    }
+
+    /**
+     * @return if the vertical linearlayout requires match_parent children remeasure
+     */
+    private boolean requiresMatchParentRemeasureForVerticalLinearLayout(int widthMeasureSpec) {
+        // HORIZONTAL measuring is handled by LinearLayout. That's why we don't need to check it
+        // here.
+        if (getOrientation() == HORIZONTAL) {
+            return false;
+        }
+
+        // When the width is not EXACT, children with MATCH_PARENT width need to be double measured.
+        // This needs to be handled in LinearLayout because NotificationOptimizedLinearLayout
+        final boolean nonExactWidth =
+                MeasureSpec.getMode(widthMeasureSpec) != MeasureSpec.EXACTLY;
+        final List<View> activeChildren = getActiveChildren();
+        for (int i = 0; i < activeChildren.size(); i++) {
+            final View child = activeChildren.get(i);
+            final ViewGroup.LayoutParams lp = child.getLayoutParams();
+            if (nonExactWidth && lp.width == ViewGroup.LayoutParams.MATCH_PARENT) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    /**
+     * @return if this layout needs to apply baseLineAlignment.
+     */
+    private boolean requiresBaselineAlignmentForHorizontalLinearLayout() {
+        // baseLineAlignment is not important for Vertical LinearLayout.
+        if (getOrientation() == VERTICAL) {
+            return false;
+        }
+        // Early return, if it is already disabled
+        if (!isBaselineAligned()) {
+            return false;
+        }
+
+        final List<View> activeChildren = getActiveChildren();
+        final int minorGravity = getGravity() & Gravity.VERTICAL_GRAVITY_MASK;
+
+        for (int i = 0; i < activeChildren.size(); i++) {
+            final View child = activeChildren.get(i);
+            if (child.getLayoutParams() instanceof LayoutParams) {
+                final LayoutParams lp = (LayoutParams) child.getLayoutParams();
+                int childBaseline = -1;
+
+                if (lp.height != LayoutParams.MATCH_PARENT) {
+                    childBaseline = child.getBaseline();
+                }
+                if (childBaseline == -1) {
+                    // This child doesn't have a baseline.
+                    continue;
+                }
+                int gravity = lp.gravity;
+                if (gravity < 0) {
+                    gravity = minorGravity;
+                }
+
+                final int result = gravity & Gravity.VERTICAL_GRAVITY_MASK;
+                if (result == Gravity.TOP || result == Gravity.BOTTOM) {
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
+
+    /**
+     * Finds the single child view within this layout that has a non-zero weight assigned to its
+     * LayoutParams.
+     *
+     * @return The weighted child view, or null if multiple weighted children exist or no weighted
+     * children are found.
+     */
+    @Nullable
+    private View getSingleWeightedChild() {
+        final boolean isVertical = getOrientation() == VERTICAL;
+        final List<View> activeChildren = getActiveChildren();
+        View singleWeightedChild = null;
+        for (int i = 0; i < activeChildren.size(); i++) {
+            final View child = activeChildren.get(i);
+            if (child.getLayoutParams() instanceof LayoutParams) {
+                final LayoutParams lp = (LayoutParams) child.getLayoutParams();
+                if ((!isVertical && lp.width == ViewGroup.LayoutParams.MATCH_PARENT)
+                        || (isVertical && lp.height == ViewGroup.LayoutParams.MATCH_PARENT)) {
+                    logSkipOptimizedOnMeasure(
+                            "There is a match parent child in the related orientation.");
+                    return null;
+                }
+                if (lp.weight != 0) {
+                    if (singleWeightedChild == null) {
+                        singleWeightedChild = child;
+                    } else {
+                        logSkipOptimizedOnMeasure("There is more than one weighted child.");
+                        return null;
+                    }
+                }
+            }
+        }
+        if (singleWeightedChild == null) {
+            logSkipOptimizedOnMeasure("There is no weighted child in this layout.");
+        } else {
+            final LayoutParams lp = (LayoutParams) singleWeightedChild.getLayoutParams();
+            boolean isHeightWrapContentOrZero =
+                    lp.height == ViewGroup.LayoutParams.WRAP_CONTENT || lp.height == 0;
+            boolean isWidthWrapContentOrZero =
+                    lp.width == ViewGroup.LayoutParams.WRAP_CONTENT || lp.width == 0;
+            if ((isVertical && !isHeightWrapContentOrZero)
+                    || (!isVertical && !isWidthWrapContentOrZero)) {
+                logSkipOptimizedOnMeasure(
+                        "Single weighted child should be either WRAP_CONTENT or 0"
+                                + " in the related orientation");
+                singleWeightedChild = null;
+            }
+        }
+
+        return singleWeightedChild;
+    }
+
+    /**
+     * Optimized measurement for the single weighted child in this LinearLayout.
+     * Measures other children, calculates remaining space, then measures the weighted
+     * child using the remaining width (or height).
+     *
+     * Note: Horizontal LinearLayout doesn't need to apply baseline in optimized case @see
+     * {@link #requiresBaselineAlignmentForHorizontalLinearLayout}.
+     *
+     * @param weightedChildView The weighted child view(with `layout_weight!=0`)
+     * @param widthMeasureSpec  The width MeasureSpec to use for measurement
+     * @param heightMeasureSpec The height MeasureSpec to use for measurement.
+     */
+    private void onMeasureOptimized(@NonNull View weightedChildView, int widthMeasureSpec,
+            int heightMeasureSpec) {
+        try {
+            if (TRACE_ONMEASURE) {
+                Trace.beginSection("NotifOptimizedLinearLayout#onMeasure");
+            }
+
+            if (getOrientation() == LinearLayout.HORIZONTAL) {
+                final ViewGroup.LayoutParams lp = weightedChildView.getLayoutParams();
+                final int childWidth = lp.width;
+                final boolean isBaselineAligned = isBaselineAligned();
+                // It should be marked 0 so that it use excessSpace in LinearLayout's onMeasure
+                lp.width = 0;
+
+                // It doesn't need to apply baseline. So disable it.
+                setBaselineAligned(false);
+                super.onMeasure(widthMeasureSpec, heightMeasureSpec);
+
+                // restore values.
+                lp.width = childWidth;
+                setBaselineAligned(isBaselineAligned);
+            } else {
+                measureVerticalOptimized(weightedChildView, widthMeasureSpec, heightMeasureSpec);
+            }
+        } finally {
+            if (TRACE_ONMEASURE) {
+                trackShouldUseOptimizedLayout();
+                Trace.endSection();
+            }
+        }
+    }
+
+    @Override
+    protected void onLayout(boolean changed, int l, int t, int r, int b) {
+        if (mShouldUseOptimizedLayout) {
+            onLayoutOptimized(changed, l, t, r, b);
+        } else {
+            super.onLayout(changed, l, t, r, b);
+        }
+    }
+
+    private void onLayoutOptimized(boolean changed, int l, int t, int r, int b) {
+        if (getOrientation() == LinearLayout.HORIZONTAL) {
+            super.onLayout(changed, l, t, r, b);
+        } else {
+            layoutVerticalOptimized(l, t, r, b);
+        }
+    }
+
+    /**
+     * Optimized measurement for the single weighted child in this LinearLayout.
+     * Measures other children, calculates remaining space, then measures the weighted
+     * child using the exact remaining height.
+     *
+     * @param weightedChildView The weighted child view(with `layout_weight=1`
+     * @param widthMeasureSpec  The width MeasureSpec to use for measurement
+     * @param heightMeasureSpec The height MeasureSpec to use for measurement.
+     */
+    private void measureVerticalOptimized(@NonNull View weightedChildView, int widthMeasureSpec,
+            int heightMeasureSpec) {
+        final int widthMode = MeasureSpec.getMode(widthMeasureSpec);
+        final int heightMode = MeasureSpec.getMode(heightMeasureSpec);
+        int maxWidth = 0;
+        int usedHeight = 0;
+        final List<View> activeChildren = getActiveChildren();
+        final int activeChildCount = activeChildren.size();
+
+        final boolean isContentFirstItem = !activeChildren.isEmpty() && activeChildren.get(0)
+                == weightedChildView;
+
+        final boolean isContentLastItem = !activeChildren.isEmpty() && activeChildren.get(
+                activeChildCount - 1) == weightedChildView;
+
+        final int horizontalPaddings = getPaddingLeft() + getPaddingRight();
+
+        // 1. Measure other child views.
+        for (int i = 0; i < activeChildCount; i++) {
+            final View child = activeChildren.get(i);
+            if (child == weightedChildView) {
+                continue;
+            }
+            final MarginLayoutParams lp = (MarginLayoutParams) child.getLayoutParams();
+
+            int requiredVerticalPadding = lp.topMargin + lp.bottomMargin;
+            if (!isContentFirstItem && i == 0) {
+                requiredVerticalPadding += getPaddingTop();
+            }
+            if (!isContentLastItem && i == activeChildCount - 1) {
+                requiredVerticalPadding += getPaddingBottom();
+            }
+
+            child.measure(ViewGroup.getChildMeasureSpec(widthMeasureSpec,
+                            horizontalPaddings + lp.leftMargin + lp.rightMargin,
+                            child.getLayoutParams().width),
+                    ViewGroup.getChildMeasureSpec(heightMeasureSpec, requiredVerticalPadding,
+                            lp.height));
+            maxWidth = Math.max(maxWidth,
+                    child.getMeasuredWidth() + lp.leftMargin + lp.rightMargin);
+            usedHeight += child.getMeasuredHeight() + requiredVerticalPadding;
+        }
+
+        // measure content
+        final MarginLayoutParams lp = (MarginLayoutParams) weightedChildView.getLayoutParams();
+
+        int usedSpace = usedHeight + lp.topMargin + lp.bottomMargin;
+        if (isContentFirstItem) {
+            usedSpace += getPaddingTop();
+        }
+        if (isContentLastItem) {
+            usedSpace += getPaddingBottom();
+        }
+
+        final int availableWidth = MeasureSpec.getSize(widthMeasureSpec);
+        final int availableHeight = MeasureSpec.getSize(heightMeasureSpec);
+
+        final int childWidthMeasureSpec = ViewGroup.getChildMeasureSpec(widthMeasureSpec,
+                horizontalPaddings + lp.leftMargin + lp.rightMargin, lp.width);
+
+        // 2. Calculate remaining height for weightedChildView.
+        final int childHeightMeasureSpec = MeasureSpec.makeMeasureSpec(
+                Math.max(0, availableHeight - usedSpace), MeasureSpec.AT_MOST);
+
+        // 3. Measure weightedChildView with the remaining remaining space.
+        weightedChildView.measure(childWidthMeasureSpec, childHeightMeasureSpec);
+        maxWidth = Math.max(maxWidth,
+                weightedChildView.getMeasuredWidth() + lp.leftMargin + lp.rightMargin);
+
+        final int totalUsedHeight = usedSpace + weightedChildView.getMeasuredHeight();
+
+        final int measuredWidth;
+        if (widthMode == MeasureSpec.EXACTLY) {
+            measuredWidth = availableWidth;
+        } else {
+            measuredWidth = maxWidth + getPaddingStart() + getPaddingEnd();
+        }
+
+        final int measuredHeight;
+        if (heightMode == MeasureSpec.EXACTLY) {
+            measuredHeight = availableHeight;
+        } else {
+            measuredHeight = totalUsedHeight;
+        }
+
+        // 4. Set the container size
+        setMeasuredDimension(
+                resolveSize(Math.max(getSuggestedMinimumWidth(), measuredWidth),
+                        widthMeasureSpec),
+                Math.max(getSuggestedMinimumHeight(), measuredHeight));
+    }
+
+    @NonNull
+    private List<View> getActiveChildren() {
+        final int childCount = getChildCount();
+        final List<View> activeChildren = new ArrayList<>();
+        for (int i = 0; i < childCount; i++) {
+            final View child = getChildAt(i);
+            if (child == null || child.getVisibility() == View.GONE) {
+                continue;
+            }
+            activeChildren.add(child);
+        }
+        return activeChildren;
+    }
+
+    //region LinearLayout copy methods
+
+    /**
+     * layoutVerticalOptimized is a version of LinearLayout's layoutVertical method that
+     * excludes
+     * TableRow-related functionalities.
+     *
+     * @see LinearLayout#onLayout(boolean, int, int, int, int)
+     */
+    private void layoutVerticalOptimized(int left, int top, int right,
+            int bottom) {
+        final int paddingLeft = mPaddingLeft;
+        final int mTotalLength = getMeasuredHeight();
+        int childTop;
+        int childLeft;
+
+        // Where right end of child should go
+        final int width = right - left;
+        int childRight = width - mPaddingRight;
+
+        // Space available for child
+        int childSpace = width - paddingLeft - mPaddingRight;
+
+        final int count = getChildCount();
+
+        final int majorGravity = getGravity() & Gravity.VERTICAL_GRAVITY_MASK;
+        final int minorGravity = getGravity() & Gravity.RELATIVE_HORIZONTAL_GRAVITY_MASK;
+
+        switch (majorGravity) {
+            case Gravity.BOTTOM:
+                // mTotalLength contains the padding already
+                childTop = mPaddingTop + bottom - top - mTotalLength;
+                break;
+
+            // mTotalLength contains the padding already
+            case Gravity.CENTER_VERTICAL:
+                childTop = mPaddingTop + (bottom - top - mTotalLength) / 2;
+                break;
+
+            case Gravity.TOP:
+            default:
+                childTop = mPaddingTop;
+                break;
+        }
+        final int dividerHeight = getDividerHeight();
+        for (int i = 0; i < count; i++) {
+            final View child = getChildAt(i);
+            if (child != null && child.getVisibility() != GONE) {
+                final int childWidth = child.getMeasuredWidth();
+                final int childHeight = child.getMeasuredHeight();
+
+                final LinearLayout.LayoutParams lp =
+                        (LinearLayout.LayoutParams) child.getLayoutParams();
+
+                int gravity = lp.gravity;
+                if (gravity < 0) {
+                    gravity = minorGravity;
+                }
+                final int layoutDirection = getLayoutDirection();
+                final int absoluteGravity = Gravity.getAbsoluteGravity(gravity,
+                        layoutDirection);
+                switch (absoluteGravity & Gravity.HORIZONTAL_GRAVITY_MASK) {
+                    case Gravity.CENTER_HORIZONTAL:
+                        childLeft =
+                                paddingLeft + ((childSpace - childWidth) / 2) + lp.leftMargin
+                                        - lp.rightMargin;
+                        break;
+
+                    case Gravity.RIGHT:
+                        childLeft = childRight - childWidth - lp.rightMargin;
+                        break;
+
+                    case Gravity.LEFT:
+                    default:
+                        childLeft = paddingLeft + lp.leftMargin;
+                        break;
+                }
+
+                if (hasDividerBeforeChildAt(i)) {
+                    childTop += dividerHeight;
+                }
+
+                childTop += lp.topMargin;
+                child.layout(childLeft, childTop, childLeft + childWidth,
+                        childTop + childHeight);
+                childTop += childHeight + lp.bottomMargin;
+
+            }
+        }
+    }
+
+    /**
+     * Used in laying out views vertically.
+     *
+     * @see #layoutVerticalOptimized
+     * @see LinearLayout#onLayout(boolean, int, int, int, int)
+     */
+    private int getDividerHeight() {
+        final Drawable dividerDrawable = getDividerDrawable();
+        if (dividerDrawable == null) {
+            return 0;
+        } else {
+            return dividerDrawable.getIntrinsicHeight();
+        }
+    }
+    //endregion
+
+    //region Logging&Tracing
+    private void trackShouldUseOptimizedLayout() {
+        if (TRACE_ONMEASURE) {
+            Trace.setCounter("NotifOptimizedLinearLayout#shouldUseOptimizedLayout",
+                    mShouldUseOptimizedLayout ? 1 : 0);
+        }
+    }
+
+    private void logSkipOptimizedOnMeasure(String reason) {
+        if (DEBUG_LAYOUT) {
+            final StringBuilder logMessage = new StringBuilder();
+            int layoutId = getId();
+            if (layoutId != NO_ID) {
+                final Resources resources = getResources();
+                if (resources != null) {
+                    logMessage.append("[");
+                    logMessage.append(resources.getResourceName(layoutId));
+                    logMessage.append("] ");
+                }
+            }
+            logMessage.append("Going to skip onMeasureOptimized reason:");
+            logMessage.append(reason);
+
+            Log.d(TAG, logMessage.toString());
+        }
+    }
+    //endregion
+}
diff --git a/core/java/com/android/internal/widget/RemeasuringLinearLayout.java b/core/java/com/android/internal/widget/RemeasuringLinearLayout.java
index 7b154a54fc85..bddad94f8688 100644
--- a/core/java/com/android/internal/widget/RemeasuringLinearLayout.java
+++ b/core/java/com/android/internal/widget/RemeasuringLinearLayout.java
@@ -30,7 +30,7 @@ import java.util.ArrayList;
  * MessagingLayouts where groups need to be able to snap it's height to.
  */
 @RemoteViews.RemoteView
-public class RemeasuringLinearLayout extends LinearLayout {
+public class RemeasuringLinearLayout extends NotificationOptimizedLinearLayout {
 
     private ArrayList<View> mMatchParentViews = new ArrayList<>();
 
@@ -79,7 +79,7 @@ public class RemeasuringLinearLayout extends LinearLayout {
             int exactHeightSpec = MeasureSpec.makeMeasureSpec(height, MeasureSpec.EXACTLY);
             for (View child : mMatchParentViews) {
                 child.measure(getChildMeasureSpec(
-                        widthMeasureSpec, getPaddingStart() + getPaddingEnd(),
+                                widthMeasureSpec, getPaddingStart() + getPaddingEnd(),
                         child.getLayoutParams().width),
                         exactHeightSpec);
             }
@@ -87,4 +87,5 @@ public class RemeasuringLinearLayout extends LinearLayout {
         mMatchParentViews.clear();
         setMeasuredDimension(getMeasuredWidth(), height);
     }
+
 }
diff --git a/core/tests/coretests/src/com/android/internal/widget/NotificationOptimizedLinearLayoutComparisonTest.java b/core/tests/coretests/src/com/android/internal/widget/NotificationOptimizedLinearLayoutComparisonTest.java
new file mode 100644
index 000000000000..08333ecd99a3
--- /dev/null
+++ b/core/tests/coretests/src/com/android/internal/widget/NotificationOptimizedLinearLayoutComparisonTest.java
@@ -0,0 +1,330 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.internal.widget;
+
+
+import static android.view.ViewGroup.LayoutParams.MATCH_PARENT;
+import static android.view.ViewGroup.LayoutParams.WRAP_CONTENT;
+
+import static org.junit.Assert.assertEquals;
+
+import android.content.Context;
+import android.platform.test.annotations.EnableFlags;
+import android.platform.test.annotations.Presubmit;
+import android.platform.test.flag.junit.SetFlagsRule;
+import android.view.Gravity;
+import android.view.View;
+import android.view.View.MeasureSpec;
+import android.view.ViewGroup;
+import android.widget.LinearLayout;
+import android.widget.flags.Flags;
+
+import androidx.test.InstrumentationRegistry;
+import androidx.test.filters.LargeTest;
+import androidx.test.runner.AndroidJUnit4;
+
+import com.google.common.truth.Expect;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Tests the consistency of {@link NotificationOptimizedLinearLayout}'s onMeasure and onLayout
+ * implementations with the behavior of the standard Android LinearLayout.
+ */
+@LargeTest
+@RunWith(AndroidJUnit4.class)
+@EnableFlags(Flags.FLAG_NOTIF_LINEARLAYOUT_OPTIMIZED)
+@Presubmit
+public class NotificationOptimizedLinearLayoutComparisonTest {
+
+    @Rule
+    public final Expect mExpect = Expect.create();
+
+    private static final int[] ORIENTATIONS = {LinearLayout.VERTICAL, LinearLayout.HORIZONTAL};
+    private static final int EXACT_SPEC = MeasureSpec.makeMeasureSpec(500,
+            MeasureSpec.EXACTLY);
+    private static final int AT_MOST_SPEC = MeasureSpec.makeMeasureSpec(500,
+            MeasureSpec.AT_MOST);
+
+    private static final int[] MEASURE_SPECS = {EXACT_SPEC, AT_MOST_SPEC};
+
+    private static final int[] GRAVITIES =
+            {Gravity.NO_GRAVITY, Gravity.TOP, Gravity.LEFT, Gravity.CENTER};
+
+    private static final int[] LAYOUT_PARAMS = {MATCH_PARENT, WRAP_CONTENT, 0, 50};
+    private static final int[] CHILD_WEIGHTS = {0, 1};
+
+    @Rule
+    public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
+
+    private Context mContext;
+
+    @Before
+    public void before() {
+        mContext = InstrumentationRegistry.getTargetContext();
+    }
+
+    @Test
+    public void test() throws Throwable {
+        for (int orientation : ORIENTATIONS) {
+            for (int widthSpec : MEASURE_SPECS) {
+                for (int heightSpec : MEASURE_SPECS) {
+                    for (int firstChildGravity : GRAVITIES) {
+                        for (int secondChildGravity : GRAVITIES) {
+                            for (int firstChildLayoutWidth : LAYOUT_PARAMS) {
+                                for (int firstChildLayoutHeight : LAYOUT_PARAMS) {
+                                    for (int secondChildLayoutWidth : LAYOUT_PARAMS) {
+                                        for (int secondChildLayoutHeight : LAYOUT_PARAMS) {
+                                            for (int firstChildWeight : CHILD_WEIGHTS) {
+                                                for (int secondChildWeight : CHILD_WEIGHTS) {
+                                                    executeTest(/*testSpec =*/createTestSpec(
+                                                            orientation,
+                                                            widthSpec, heightSpec,
+                                                            firstChildLayoutWidth,
+                                                            firstChildLayoutHeight,
+                                                            secondChildLayoutWidth,
+                                                            secondChildLayoutHeight,
+                                                            firstChildGravity,
+                                                            secondChildGravity,
+                                                            firstChildWeight,
+                                                            secondChildWeight));
+                                                }
+                                            }
+                                        }
+                                    }
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    private void executeTest(TestSpec testSpec) {
+        // GIVEN
+        final List<View> controlChildren =
+                new ArrayList<>();
+        final List<View> testChildren =
+                new ArrayList<>();
+
+        controlChildren.add(
+                buildChildView(
+                        testSpec.mFirstChildLayoutWidth,
+                        testSpec.mFirstChildLayoutHeight,
+                        testSpec.mFirstChildGravity,
+                        testSpec.mFirstChildWeight));
+        controlChildren.add(
+                buildChildView(
+                        testSpec.mSecondChildLayoutWidth,
+                        testSpec.mSecondChildLayoutHeight,
+                        testSpec.mSecondChildGravity,
+                        testSpec.mSecondChildWeight));
+
+        testChildren.add(
+                buildChildView(
+                        testSpec.mFirstChildLayoutWidth,
+                        testSpec.mFirstChildLayoutHeight,
+                        testSpec.mFirstChildGravity,
+                        testSpec.mFirstChildWeight));
+        testChildren.add(
+                buildChildView(
+                        testSpec.mSecondChildLayoutWidth,
+                        testSpec.mSecondChildLayoutHeight,
+                        testSpec.mSecondChildGravity,
+                        testSpec.mSecondChildWeight));
+
+        final LinearLayout controlContainer = buildLayout(false,
+                testSpec.mOrientation,
+                controlChildren);
+
+        final LinearLayout testContainer = buildLayout(true,
+                testSpec.mOrientation,
+                testChildren);
+
+        // WHEN
+        controlContainer.measure(testSpec.mWidthSpec, testSpec.mHeightSpec);
+        testContainer.measure(testSpec.mWidthSpec, testSpec.mHeightSpec);
+        controlContainer.layout(0, 0, 1000, 1000);
+        testContainer.layout(0, 0, 1000, 1000);
+        // THEN
+        assertLayoutsEqual("Test Case:" + testSpec, controlContainer, testContainer);
+    }
+
+    private static class TestSpec {
+        private final int mOrientation;
+        private final int mWidthSpec;
+        private final int mHeightSpec;
+        private final int mFirstChildLayoutWidth;
+        private final int mFirstChildLayoutHeight;
+        private final int mSecondChildLayoutWidth;
+        private final int mSecondChildLayoutHeight;
+        private final int mFirstChildGravity;
+        private final int mSecondChildGravity;
+        private final int mFirstChildWeight;
+        private final int mSecondChildWeight;
+
+        TestSpec(
+                int orientation,
+                int widthSpec,
+                int heightSpec,
+                int firstChildLayoutWidth,
+                int firstChildLayoutHeight,
+                int secondChildLayoutWidth,
+                int secondChildLayoutHeight,
+                int firstChildGravity,
+                int secondChildGravity,
+                int firstChildWeight,
+                int secondChildWeight) {
+            mOrientation = orientation;
+            mWidthSpec = widthSpec;
+            mHeightSpec = heightSpec;
+            mFirstChildLayoutWidth = firstChildLayoutWidth;
+            mFirstChildLayoutHeight = firstChildLayoutHeight;
+            mSecondChildLayoutWidth = secondChildLayoutWidth;
+            mSecondChildLayoutHeight = secondChildLayoutHeight;
+            mFirstChildGravity = firstChildGravity;
+            mSecondChildGravity = secondChildGravity;
+            mFirstChildWeight = firstChildWeight;
+            mSecondChildWeight = secondChildWeight;
+        }
+
+        @Override
+        public String toString() {
+            return "TestSpec{"
+                    + "mOrientation=" + orientationToString(mOrientation)
+                    + ", mWidthSpec=" + MeasureSpec.toString(mWidthSpec)
+                    + ", mHeightSpec=" + MeasureSpec.toString(mHeightSpec)
+                    + ", mFirstChildLayoutWidth=" + sizeToString(mFirstChildLayoutWidth)
+                    + ", mFirstChildLayoutHeight=" + sizeToString(mFirstChildLayoutHeight)
+                    + ", mSecondChildLayoutWidth=" + sizeToString(mSecondChildLayoutWidth)
+                    + ", mSecondChildLayoutHeight=" + sizeToString(mSecondChildLayoutHeight)
+                    + ", mFirstChildGravity=" + mFirstChildGravity
+                    + ", mSecondChildGravity=" + mSecondChildGravity
+                    + ", mFirstChildWeight=" + mFirstChildWeight
+                    + ", mSecondChildWeight=" + mSecondChildWeight
+                    + '}';
+        }
+
+        private String orientationToString(int orientation) {
+            if (orientation == LinearLayout.VERTICAL) {
+                return "vertical";
+            } else if (orientation == LinearLayout.HORIZONTAL) {
+                return "horizontal";
+            }
+            throw new IllegalArgumentException();
+        }
+
+        private String sizeToString(int size) {
+            if (size == WRAP_CONTENT) {
+                return "wrap-content";
+            }
+            if (size == MATCH_PARENT) {
+                return "match-parent";
+            }
+            return String.valueOf(size);
+        }
+    }
+
+    private LinearLayout buildLayout(boolean isNotificationOptimized,
+            @LinearLayout.OrientationMode int orientation, List<View> children) {
+        final LinearLayout linearLayout;
+        if (isNotificationOptimized) {
+            linearLayout = new NotificationOptimizedLinearLayout(mContext);
+        } else {
+            linearLayout = new LinearLayout(mContext);
+        }
+        linearLayout.setOrientation(orientation);
+        for (int i = 0; i < children.size(); i++) {
+            linearLayout.addView(children.get(i));
+        }
+        return linearLayout;
+    }
+
+    private void assertLayoutsEqual(String testCase, View controlView, View testView) {
+        mExpect.withMessage("MeasuredWidths are not equal. Test Case:" + testCase)
+                .that(testView.getMeasuredWidth()).isEqualTo(controlView.getMeasuredWidth());
+        mExpect.withMessage("MeasuredHeights are not equal. Test Case:" + testCase)
+                .that(testView.getMeasuredHeight()).isEqualTo(controlView.getMeasuredHeight());
+        mExpect.withMessage("Left Positions are not equal. Test Case:" + testCase)
+                .that(testView.getLeft()).isEqualTo(controlView.getLeft());
+        mExpect.withMessage("Top Positions are not equal. Test Case:" + testCase)
+                .that(testView.getTop()).isEqualTo(controlView.getTop());
+        if (controlView instanceof ViewGroup && testView instanceof ViewGroup) {
+            final ViewGroup controlGroup = (ViewGroup) controlView;
+            final ViewGroup testGroup = (ViewGroup) testView;
+            // Test and Control Views should be identical by hierarchy for the comparison.
+            // That's why mExpect is not used here for assertion.
+            assertEquals(controlGroup.getChildCount(), testGroup.getChildCount());
+
+            for (int i = 0; i < controlGroup.getChildCount(); i++) {
+                View controlChild = controlGroup.getChildAt(i);
+                View testChild = testGroup.getChildAt(i);
+
+                assertLayoutsEqual(testCase, controlChild, testChild);
+            }
+        }
+    }
+
+    private static class TestView extends View {
+        TestView(Context context) {
+            super(context);
+        }
+
+        @Override
+        public int getBaseline() {
+            return 5;
+        }
+    }
+
+
+    private TestSpec createTestSpec(int orientation,
+            int widthSpec, int heightSpec,
+            int firstChildLayoutWidth, int firstChildLayoutHeight, int secondChildLayoutWidth,
+            int secondChildLayoutHeight, int firstChildGravity, int secondChildGravity,
+            int firstChildWeight, int secondChildWeight) {
+
+        return new TestSpec(
+                orientation,
+                widthSpec, heightSpec,
+                firstChildLayoutWidth,
+                firstChildLayoutHeight,
+                secondChildLayoutWidth,
+                secondChildLayoutHeight,
+                firstChildGravity,
+                secondChildGravity,
+                firstChildWeight,
+                secondChildWeight);
+    }
+
+    private View buildChildView(int childLayoutWidth, int childLayoutHeight,
+            int childGravity, int childWeight) {
+        final View childView = new TestView(mContext);
+        // Set desired size using LayoutParams
+        LinearLayout.LayoutParams params = new LinearLayout.LayoutParams(childLayoutWidth,
+                childLayoutHeight, childWeight);
+        params.gravity = childGravity;
+        childView.setLayoutParams(params);
+        return childView;
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotifRemoteViewsFactoryContainer.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotifRemoteViewsFactoryContainer.kt
index 99177c270b32..195fe785b538 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotifRemoteViewsFactoryContainer.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotifRemoteViewsFactoryContainer.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.statusbar.notification.row
 
+import android.widget.flags.Flags.notifLinearlayoutOptimized
 import com.android.systemui.flags.FeatureFlags
 import com.android.systemui.flags.Flags
 import javax.inject.Inject
@@ -31,6 +32,7 @@ constructor(
     precomputedTextViewFactory: PrecomputedTextViewFactory,
     bigPictureLayoutInflaterFactory: BigPictureLayoutInflaterFactory,
     callLayoutSetDataAsyncFactory: CallLayoutSetDataAsyncFactory,
+    optimizedLinearLayoutFactory: NotificationOptimizedLinearLayoutFactory
 ) : NotifRemoteViewsFactoryContainer {
     override val factories: Set<NotifRemoteViewsFactory> = buildSet {
         add(precomputedTextViewFactory)
@@ -40,5 +42,8 @@ constructor(
         if (featureFlags.isEnabled(Flags.CALL_LAYOUT_ASYNC_SET_DATA)) {
             add(callLayoutSetDataAsyncFactory)
         }
+        if (notifLinearlayoutOptimized()) {
+            add(optimizedLinearLayoutFactory)
+        }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationOptimizedLinearLayoutFactory.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationOptimizedLinearLayoutFactory.kt
new file mode 100644
index 000000000000..f231fbc6e433
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationOptimizedLinearLayoutFactory.kt
@@ -0,0 +1,41 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.statusbar.notification.row
+
+import android.content.Context
+import android.util.AttributeSet
+import android.view.View
+import android.widget.LinearLayout
+import com.android.internal.widget.NotificationOptimizedLinearLayout
+import javax.inject.Inject
+
+class NotificationOptimizedLinearLayoutFactory @Inject constructor() : NotifRemoteViewsFactory {
+    override fun instantiate(
+        row: ExpandableNotificationRow,
+        @NotificationRowContentBinder.InflationFlag layoutType: Int,
+        parent: View?,
+        name: String,
+        context: Context,
+        attrs: AttributeSet
+    ): View? {
+        return when (name) {
+            LinearLayout::class.java.name,
+            LinearLayout::class.java.simpleName -> NotificationOptimizedLinearLayout(context, attrs)
+            else -> null
+        }
+    }
+}