diff --git a/libs/WindowManager/Shell/Android.bp b/libs/WindowManager/Shell/Android.bp
index cdff585..e9b3c49 100644
--- a/libs/WindowManager/Shell/Android.bp
+++ b/libs/WindowManager/Shell/Android.bp
@@ -43,7 +43,7 @@
     name: "wm_shell_util-sources",
     srcs: [
         "src/com/android/wm/shell/util/**/*.java",
-        "src/com/android/wm/shell/common/split/SplitScreenConstants.java"
+        "src/com/android/wm/shell/common/split/SplitScreenConstants.java",
     ],
     path: "src",
 }
@@ -74,13 +74,13 @@
     ],
     tools: ["protologtool"],
     cmd: "$(location protologtool) transform-protolog-calls " +
-      "--protolog-class com.android.internal.protolog.common.ProtoLog " +
-      "--protolog-impl-class com.android.wm.shell.protolog.ShellProtoLogImpl " +
-      "--protolog-cache-class com.android.wm.shell.protolog.ShellProtoLogCache " +
-      "--loggroups-class com.android.wm.shell.protolog.ShellProtoLogGroup " +
-      "--loggroups-jar $(location :wm_shell_protolog-groups) " +
-      "--output-srcjar $(out) " +
-      "$(locations :wm_shell-sources)",
+        "--protolog-class com.android.internal.protolog.common.ProtoLog " +
+        "--protolog-impl-class com.android.wm.shell.protolog.ShellProtoLogImpl " +
+        "--protolog-cache-class com.android.wm.shell.protolog.ShellProtoLogCache " +
+        "--loggroups-class com.android.wm.shell.protolog.ShellProtoLogGroup " +
+        "--loggroups-jar $(location :wm_shell_protolog-groups) " +
+        "--output-srcjar $(out) " +
+        "$(locations :wm_shell-sources)",
     out: ["wm_shell_protolog.srcjar"],
 }
 
@@ -92,13 +92,14 @@
     ],
     tools: ["protologtool"],
     cmd: "$(location protologtool) generate-viewer-config " +
-      "--protolog-class com.android.internal.protolog.common.ProtoLog " +
-      "--loggroups-class com.android.wm.shell.protolog.ShellProtoLogGroup " +
-      "--loggroups-jar $(location :wm_shell_protolog-groups) " +
-      "--viewer-conf $(out) " +
-      "$(locations :wm_shell-sources)",
+        "--protolog-class com.android.internal.protolog.common.ProtoLog " +
+        "--loggroups-class com.android.wm.shell.protolog.ShellProtoLogGroup " +
+        "--loggroups-jar $(location :wm_shell_protolog-groups) " +
+        "--viewer-conf $(out) " +
+        "$(locations :wm_shell-sources)",
     out: ["wm_shell_protolog.json"],
 }
+
 // End ProtoLog
 
 java_library {
@@ -123,11 +124,12 @@
         "res",
     ],
     java_resources: [
-        ":generate-wm_shell_protolog.json"
+        ":generate-wm_shell_protolog.json",
     ],
     static_libs: [
         "androidx.appcompat_appcompat",
         "androidx.arch.core_core-runtime",
+        "androidx-constraintlayout_constraintlayout",
         "androidx.dynamicanimation_dynamicanimation",
         "androidx.recyclerview_recyclerview",
         "kotlinx-coroutines-android",
diff --git a/libs/WindowManager/Shell/res/layout/badged_image_view.xml b/libs/WindowManager/Shell/res/layout/badged_image_view.xml
new file mode 100644
index 0000000..5f07121
--- /dev/null
+++ b/libs/WindowManager/Shell/res/layout/badged_image_view.xml
@@ -0,0 +1,55 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  ~ Copyright (C) 2021 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.
+  -->
+<merge xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:app="http://schemas.android.com/apk/res-auto">
+
+    <ImageView
+        android:id="@+id/icon_view"
+        android:layout_width="match_parent"
+        android:layout_height="match_parent"
+        android:contentDescription="@null" />
+
+    <!--
+        Icon badge size is defined in Launcher3 BaseIconFactory as 0.444 of icon size.
+        Constraint guide starts from left, which means for a badge positioned on the right,
+        percent has to be 1 - 0.444 to have the same effect.
+    -->
+    <androidx.constraintlayout.widget.Guideline
+        android:id="@+id/app_icon_constraint_horizontal"
+        android:layout_width="wrap_content"
+        android:layout_height="wrap_content"
+        android:orientation="horizontal"
+        app:layout_constraintGuide_percent="0.556" />
+
+    <androidx.constraintlayout.widget.Guideline
+        android:id="@+id/app_icon_constraint_vertical"
+        android:layout_width="wrap_content"
+        android:layout_height="wrap_content"
+        android:orientation="vertical"
+        app:layout_constraintGuide_percent="0.556" />
+
+    <ImageView
+        android:id="@+id/app_icon_view"
+        android:layout_width="0dp"
+        android:layout_height="0dp"
+        android:contentDescription="@null"
+        app:layout_constraintBottom_toBottomOf="parent"
+        app:layout_constraintLeft_toLeftOf="@id/app_icon_constraint_vertical"
+        app:layout_constraintRight_toRightOf="parent"
+        app:layout_constraintTop_toTopOf="@id/app_icon_constraint_horizontal" />
+
+</merge>
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BadgedImageView.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BadgedImageView.java
index 686fbbf..c52d87d 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BadgedImageView.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BadgedImageView.java
@@ -19,7 +19,6 @@
 import android.annotation.Nullable;
 import android.content.Context;
 import android.content.res.TypedArray;
-import android.graphics.Bitmap;
 import android.graphics.Canvas;
 import android.graphics.Outline;
 import android.graphics.Path;
@@ -27,14 +26,16 @@
 import android.graphics.drawable.Drawable;
 import android.util.AttributeSet;
 import android.util.PathParser;
-import android.view.Gravity;
+import android.view.LayoutInflater;
 import android.view.View;
 import android.view.ViewOutlineProvider;
-import android.widget.FrameLayout;
 import android.widget.ImageView;
 
+import androidx.constraintlayout.widget.ConstraintLayout;
+
 import com.android.launcher3.icons.DotRenderer;
 import com.android.launcher3.icons.IconNormalizer;
+import com.android.wm.shell.R;
 import com.android.wm.shell.animation.Interpolators;
 
 import java.util.EnumSet;
@@ -46,14 +47,12 @@
  * Badge = the icon associated with the app that created this bubble, this will show work profile
  * badge if appropriate.
  */
-public class BadgedImageView extends FrameLayout {
+public class BadgedImageView extends ConstraintLayout {
 
     /** Same value as Launcher3 dot code */
     public static final float WHITE_SCRIM_ALPHA = 0.54f;
     /** Same as value in Launcher3 IconShape */
     public static final int DEFAULT_PATH_SIZE = 100;
-    /** Same as value in Launcher3 BaseIconFactory */
-    private static final float ICON_BADGE_SCALE = 0.444f;
 
     /**
      * Flags that suppress the visibility of the 'new' dot, for one reason or another. If any of
@@ -105,11 +104,13 @@
     public BadgedImageView(Context context, AttributeSet attrs, int defStyleAttr,
             int defStyleRes) {
         super(context, attrs, defStyleAttr, defStyleRes);
+        // We manage positioning the badge ourselves
+        setLayoutDirection(LAYOUT_DIRECTION_LTR);
 
-        mBubbleIcon = new ImageView(context);
-        addView(mBubbleIcon);
-        mAppIcon = new ImageView(context);
-        addView(mAppIcon);
+        LayoutInflater.from(context).inflate(R.layout.badged_image_view, this);
+
+        mBubbleIcon = findViewById(R.id.icon_view);
+        mAppIcon = findViewById(R.id.app_icon_view);
 
         final TypedArray ta = mContext.obtainStyledAttributes(attrs, new int[]{android.R.attr.src},
                 defStyleAttr, defStyleRes);
@@ -161,6 +162,7 @@
     public void setRenderedBubble(BubbleViewProvider bubble) {
         mBubble = bubble;
         mBubbleIcon.setImageBitmap(bubble.getBubbleIcon());
+        mAppIcon.setImageBitmap(bubble.getAppBadge());
         if (mDotSuppressionFlags.contains(SuppressionFlag.BEHIND_STACK)) {
             hideBadge();
         } else {
@@ -348,26 +350,17 @@
     }
 
     void showBadge() {
-        Bitmap badge = mBubble.getAppBadge();
-        if (badge == null) {
+        if (mBubble.getAppBadge() == null) {
             mAppIcon.setVisibility(GONE);
             return;
         }
-
-        final int bubbleSize = mBubble.getBubbleIcon().getWidth();
-        final int badgeSize = (int) (ICON_BADGE_SCALE * bubbleSize);
-
-        FrameLayout.LayoutParams appIconParams = (LayoutParams) mAppIcon.getLayoutParams();
-        appIconParams.height = badgeSize;
-        appIconParams.width = badgeSize;
+        int translationX;
         if (mOnLeft) {
-            appIconParams.gravity = Gravity.BOTTOM | Gravity.LEFT;
+            translationX = -(mBubbleIcon.getWidth() - mAppIcon.getWidth());
         } else {
-            appIconParams.gravity = Gravity.BOTTOM | Gravity.RIGHT;
+            translationX = 0;
         }
-        mAppIcon.setLayoutParams(appIconParams);
-
-        mAppIcon.setImageBitmap(badge);
+        mAppIcon.setTranslationX(translationX);
         mAppIcon.setVisibility(VISIBLE);
     }