diff --git a/core/java/android/provider/Settings.java b/core/java/android/provider/Settings.java
index 589ecc0a4de50..dc4a74766ed92 100644
--- a/core/java/android/provider/Settings.java
+++ b/core/java/android/provider/Settings.java
@@ -10119,6 +10119,61 @@ public final class Settings {
         @Readable
         public static final String GAME_DASHBOARD_ALWAYS_ON = "game_dashboard_always_on";
 
+
+        /**
+         * For this device state, no specific auto-rotation lock setting should be applied.
+         * If the user toggles the auto-rotate lock in this state, the setting will apply to the
+         * previously valid device state.
+         * @hide
+         */
+        public static final int DEVICE_STATE_ROTATION_LOCK_IGNORED = 0;
+        /**
+         * For this device state, the setting for auto-rotation is locked.
+         * @hide
+         */
+        public static final int DEVICE_STATE_ROTATION_LOCK_LOCKED = 1;
+        /**
+         * For this device state, the setting for auto-rotation is unlocked.
+         * @hide
+         */
+        public static final int DEVICE_STATE_ROTATION_LOCK_UNLOCKED = 2;
+
+        /**
+         * The different settings that can be used as values with
+         * {@link #DEVICE_STATE_ROTATION_LOCK}.
+         * @hide
+         */
+        @IntDef(prefix = {"DEVICE_STATE_ROTATION_LOCK_"}, value = {
+                DEVICE_STATE_ROTATION_LOCK_IGNORED,
+                DEVICE_STATE_ROTATION_LOCK_LOCKED,
+                DEVICE_STATE_ROTATION_LOCK_UNLOCKED,
+        })
+        @Retention(RetentionPolicy.SOURCE)
+        @interface DeviceStateRotationLockSetting {
+        }
+
+        /**
+         * Rotation lock setting keyed on device state.
+         *
+         * This holds a serialized map using int keys that represent Device States and value of
+         * {@link DeviceStateRotationLockSetting} representing the rotation lock setting for that
+         * device state.
+         *
+         * Serialized as key0:value0:key1:value1:...:keyN:valueN.
+         *
+         * Example: "0:1:1:2:2:1"
+         * This example represents a map of:
+         * <ul>
+         *     <li>0 -> DEVICE_STATE_ROTATION_LOCK_LOCKED</li>
+         *     <li>1 -> DEVICE_STATE_ROTATION_LOCK_UNLOCKED</li>
+         *     <li>2 -> DEVICE_STATE_ROTATION_LOCK_IGNORED</li>
+         * </ul>
+         *
+         * @hide
+         */
+        public static final String DEVICE_STATE_ROTATION_LOCK =
+                "device_state_rotation_lock";
+
         /**
          * These entries are considered common between the personal and the managed profile,
          * since the managed profile doesn't get to change them.
diff --git a/core/res/res/values/config.xml b/core/res/res/values/config.xml
index 1586f164a0a5d..a3099788a46e9 100644
--- a/core/res/res/values/config.xml
+++ b/core/res/res/values/config.xml
@@ -666,6 +666,17 @@
          display is powered on at the same time. -->
     <bool name="config_supportsConcurrentInternalDisplays">true</bool>
 
+    <!-- Map of DeviceState to rotation lock setting. Each entry must be in the format
+         "key:value", for example: "0:1".
+          The keys are device states, and the values are one of
+          Settings.Secure.DeviceStateRotationLockSetting.
+          Any device state that doesn't have a default set here will be treated as
+          DEVICE_STATE_ROTATION_LOCK_IGNORED meaning it will not have its own rotation lock setting.
+          If this map is missing, the feature is disabled and only one global rotation lock setting
+           will apply, regardless of device state. -->
+    <string-array name="config_perDeviceStateRotationLockDefaults" />
+
+
     <!-- Desk dock behavior -->
 
     <!-- The number of degrees to rotate the display when the device is in a desk dock.
diff --git a/core/res/res/values/symbols.xml b/core/res/res/values/symbols.xml
index cd590cbbe2b0c..7c96d9dfb1c30 100644
--- a/core/res/res/values/symbols.xml
+++ b/core/res/res/values/symbols.xml
@@ -3838,6 +3838,8 @@
   <java-symbol type="string" name="config_foldedArea" />
   <java-symbol type="bool" name="config_supportsConcurrentInternalDisplays" />
   <java-symbol type="bool" name="config_unfoldTransitionEnabled" />
+  <java-symbol type="array" name="config_perDeviceStateRotationLockDefaults" />
+
 
   <java-symbol type="array" name="config_disableApksUnlessMatchedSku_apk_list" />
   <java-symbol type="array" name="config_disableApkUnlessMatchedSku_skus_list" />
diff --git a/packages/SettingsProvider/src/android/provider/settings/validators/SecureSettingsValidators.java b/packages/SettingsProvider/src/android/provider/settings/validators/SecureSettingsValidators.java
index 7a9f81e088cc1..7aeacdc0cf71c 100644
--- a/packages/SettingsProvider/src/android/provider/settings/validators/SecureSettingsValidators.java
+++ b/packages/SettingsProvider/src/android/provider/settings/validators/SecureSettingsValidators.java
@@ -34,7 +34,9 @@ import static android.provider.settings.validators.SettingsValidators.TILE_LIST_
 import static android.provider.settings.validators.SettingsValidators.TTS_LIST_VALIDATOR;
 
 import android.provider.Settings.Secure;
+import android.text.TextUtils;
 import android.util.ArrayMap;
+import android.util.ArraySet;
 
 import java.util.Map;
 
@@ -287,5 +289,32 @@ public class SecureSettingsValidators {
         VALIDATORS.put(Secure.CLIPBOARD_SHOW_ACCESS_NOTIFICATIONS, BOOLEAN_VALIDATOR);
         VALIDATORS.put(Secure.NOTIFICATION_BUBBLES, BOOLEAN_VALIDATOR);
         VALIDATORS.put(Secure.LOCATION_TIME_ZONE_DETECTION_ENABLED, BOOLEAN_VALIDATOR);
+        VALIDATORS.put(Secure.DEVICE_STATE_ROTATION_LOCK, value -> {
+            if (TextUtils.isEmpty(value)) {
+                return true;
+            }
+            String[] intValues = value.split(":");
+            if (intValues.length % 2 != 0) {
+                return false;
+            }
+            InclusiveIntegerRangeValidator enumValidator =
+                    new InclusiveIntegerRangeValidator(
+                            Secure.DEVICE_STATE_ROTATION_LOCK_IGNORED,
+                            Secure.DEVICE_STATE_ROTATION_LOCK_UNLOCKED);
+            ArraySet<String> keys = new ArraySet<>();
+            for (int i = 0; i < intValues.length - 1; ) {
+                String entryKey = intValues[i++];
+                String entryValue = intValues[i++];
+                if (!NON_NEGATIVE_INTEGER_VALIDATOR.validate(entryKey)
+                        || !enumValidator.validate(entryValue)) {
+                    return false;
+                }
+                // If the same device state key was specified more than once, this is invalid
+                if (!keys.add(entryKey)) {
+                    return false;
+                }
+            }
+            return true;
+        });
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/DeviceStateRotationLockSettingController.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/DeviceStateRotationLockSettingController.java
new file mode 100644
index 0000000000000..41cacf5142fdd
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/DeviceStateRotationLockSettingController.java
@@ -0,0 +1,239 @@
+/*
+ * 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.
+ */
+
+package com.android.systemui.statusbar.policy;
+
+
+import static android.provider.Settings.Secure.DEVICE_STATE_ROTATION_LOCK_IGNORED;
+import static android.provider.Settings.Secure.DEVICE_STATE_ROTATION_LOCK_LOCKED;
+import static android.provider.Settings.Secure.DEVICE_STATE_ROTATION_LOCK_UNLOCKED;
+
+import static com.android.systemui.statusbar.policy.dagger.StatusBarPolicyModule.DEVICE_STATE_ROTATION_LOCK_DEFAULTS;
+
+import android.annotation.Nullable;
+import android.hardware.devicestate.DeviceStateManager;
+import android.os.UserHandle;
+import android.provider.Settings;
+import android.text.TextUtils;
+import android.util.Log;
+import android.util.SparseIntArray;
+
+import com.android.systemui.dagger.SysUISingleton;
+import com.android.systemui.dagger.qualifiers.Main;
+import com.android.systemui.util.settings.SecureSettings;
+import com.android.systemui.util.wrapper.RotationPolicyWrapper;
+
+import java.util.concurrent.Executor;
+
+import javax.inject.Inject;
+import javax.inject.Named;
+
+/**
+ * Handles reading and writing of rotation lock settings per device state, as well as setting
+ * the rotation lock when device state changes.
+ **/
+@SysUISingleton
+public final class DeviceStateRotationLockSettingController implements Listenable,
+        RotationLockController.RotationLockControllerCallback {
+
+    private static final String TAG = "DSRotateLockSettingCon";
+
+    private static final String SEPARATOR_REGEX = ":";
+
+    private final SecureSettings mSecureSettings;
+    private final RotationPolicyWrapper mRotationPolicyWrapper;
+    private final DeviceStateManager mDeviceStateManager;
+    private final Executor mMainExecutor;
+    private final String[] mDeviceStateRotationLockDefaults;
+
+    private SparseIntArray mDeviceStateRotationLockSettings;
+    // TODO(b/183001527): Add API to query current device state and initialize this.
+    private int mDeviceState = -1;
+    @Nullable
+    private DeviceStateManager.DeviceStateCallback mDeviceStateCallback;
+
+
+    @Inject
+    public DeviceStateRotationLockSettingController(
+            SecureSettings secureSettings,
+            RotationPolicyWrapper rotationPolicyWrapper,
+            DeviceStateManager deviceStateManager,
+            @Main Executor executor,
+            @Named(DEVICE_STATE_ROTATION_LOCK_DEFAULTS) String[] deviceStateRotationLockDefaults
+    ) {
+        mSecureSettings = secureSettings;
+        mRotationPolicyWrapper = rotationPolicyWrapper;
+        mDeviceStateManager = deviceStateManager;
+        mMainExecutor = executor;
+        mDeviceStateRotationLockDefaults = deviceStateRotationLockDefaults;
+    }
+
+    /**
+     * Loads the settings from storage.
+     */
+    public void initialize() {
+        String serializedSetting =
+                mSecureSettings.getStringForUser(Settings.Secure.DEVICE_STATE_ROTATION_LOCK,
+                        UserHandle.USER_CURRENT);
+        if (TextUtils.isEmpty(serializedSetting)) {
+            // No settings saved, we should load the defaults and persist them.
+            fallbackOnDefaults();
+            return;
+        }
+        String[] values = serializedSetting.split(SEPARATOR_REGEX);
+        if (values.length % 2 != 0) {
+            // Each entry should be a key/value pair, so this is corrupt.
+            Log.wtf(TAG, "Can't deserialize saved settings, falling back on defaults");
+            fallbackOnDefaults();
+            return;
+        }
+        mDeviceStateRotationLockSettings = new SparseIntArray(values.length / 2);
+        int key;
+        int value;
+
+        for (int i = 0; i < values.length - 1; ) {
+            try {
+                key = Integer.parseInt(values[i++]);
+                value = Integer.parseInt(values[i++]);
+                mDeviceStateRotationLockSettings.put(key, value);
+            } catch (NumberFormatException e) {
+                Log.wtf(TAG, "Error deserializing one of the saved settings", e);
+                fallbackOnDefaults();
+                return;
+            }
+        }
+    }
+
+    private void fallbackOnDefaults() {
+        loadDefaults();
+        persistSettings();
+    }
+
+    @Override
+    public void setListening(boolean listening) {
+        if (listening) {
+            // Note that this is called once with the initial state of the device, even if there
+            // is no user action.
+            mDeviceStateCallback = this::updateDeviceState;
+            mDeviceStateManager.registerCallback(mMainExecutor, mDeviceStateCallback);
+        } else {
+            if (mDeviceStateCallback != null) {
+                mDeviceStateManager.unregisterCallback(mDeviceStateCallback);
+            }
+        }
+    }
+
+    @Override
+    public void onRotationLockStateChanged(boolean rotationLocked, boolean affordanceVisible) {
+        if (mDeviceState == -1) {
+            Log.wtf(TAG, "Device state was not initialized.");
+            return;
+        }
+
+        if (rotationLocked == isRotationLockedForCurrentState()) {
+            Log.v(TAG, "Rotation lock same as the current setting, no need to update.");
+            return;
+        }
+
+        saveNewRotationLockSetting(rotationLocked);
+    }
+
+    private void saveNewRotationLockSetting(boolean isRotationLocked) {
+        Log.v(TAG, "saveNewRotationLockSetting [state=" + mDeviceState + "] [isRotationLocked="
+                + isRotationLocked + "]");
+
+        mDeviceStateRotationLockSettings.put(mDeviceState,
+                isRotationLocked
+                        ? DEVICE_STATE_ROTATION_LOCK_LOCKED
+                        : DEVICE_STATE_ROTATION_LOCK_UNLOCKED);
+        persistSettings();
+    }
+
+    private boolean isRotationLockedForCurrentState() {
+        return mDeviceStateRotationLockSettings.get(mDeviceState,
+                DEVICE_STATE_ROTATION_LOCK_IGNORED) == DEVICE_STATE_ROTATION_LOCK_LOCKED;
+    }
+
+    private void updateDeviceState(int state) {
+        Log.v(TAG, "updateDeviceState [state=" + state + "]");
+        if (mDeviceState == state) {
+            return;
+        }
+
+        int rotationLockSetting =
+                mDeviceStateRotationLockSettings.get(state, DEVICE_STATE_ROTATION_LOCK_IGNORED);
+        if (rotationLockSetting == DEVICE_STATE_ROTATION_LOCK_IGNORED) {
+            // We won't handle this device state. The same rotation lock setting as before should
+            // apply and any changes to the rotation lock setting will be written for the previous
+            // valid device state.
+            Log.v(TAG, "Ignoring new device state: " + state);
+            return;
+        }
+
+        // Accept the new state
+        mDeviceState = state;
+
+        // Update the rotation lock setting if needed for this new device state
+        boolean newRotationLockSetting = rotationLockSetting == DEVICE_STATE_ROTATION_LOCK_LOCKED;
+        if (newRotationLockSetting != mRotationPolicyWrapper.isRotationLocked()) {
+            mRotationPolicyWrapper.setRotationLock(newRotationLockSetting);
+        }
+    }
+
+    private void persistSettings() {
+        if (mDeviceStateRotationLockSettings.size() == 0) {
+            mSecureSettings.putStringForUser(Settings.Secure.DEVICE_STATE_ROTATION_LOCK,
+                    /* value= */"", UserHandle.USER_CURRENT);
+            return;
+        }
+
+        StringBuilder stringBuilder = new StringBuilder();
+        stringBuilder.append(mDeviceStateRotationLockSettings.keyAt(0))
+                .append(SEPARATOR_REGEX)
+                .append(mDeviceStateRotationLockSettings.valueAt(0));
+
+        for (int i = 1; i < mDeviceStateRotationLockSettings.size(); i++) {
+            stringBuilder
+                    .append(SEPARATOR_REGEX)
+                    .append(mDeviceStateRotationLockSettings.keyAt(i))
+                    .append(SEPARATOR_REGEX)
+                    .append(mDeviceStateRotationLockSettings.valueAt(i));
+        }
+        mSecureSettings.putStringForUser(Settings.Secure.DEVICE_STATE_ROTATION_LOCK,
+                stringBuilder.toString(), UserHandle.USER_CURRENT);
+    }
+
+    private void loadDefaults() {
+        if (mDeviceStateRotationLockDefaults.length == 0) {
+            Log.w(TAG, "Empty default settings");
+            mDeviceStateRotationLockSettings = new SparseIntArray(/* initialCapacity= */0);
+            return;
+        }
+        mDeviceStateRotationLockSettings =
+                new SparseIntArray(mDeviceStateRotationLockDefaults.length);
+        for (String serializedDefault : mDeviceStateRotationLockDefaults) {
+            String[] entry = serializedDefault.split(SEPARATOR_REGEX);
+            try {
+                int key = Integer.parseInt(entry[0]);
+                int value = Integer.parseInt(entry[1]);
+                mDeviceStateRotationLockSettings.put(key, value);
+            } catch (NumberFormatException e) {
+                Log.wtf(TAG, "Error deserializing default settings", e);
+            }
+        }
+    }
+
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/RotationLockControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/RotationLockControllerImpl.java
index 53d68d0ff0ac4..67f5364e3d3ff 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/RotationLockControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/RotationLockControllerImpl.java
@@ -16,36 +16,54 @@
 
 package com.android.systemui.statusbar.policy;
 
-import android.content.Context;
+import static com.android.systemui.statusbar.policy.dagger.StatusBarPolicyModule.DEVICE_STATE_ROTATION_LOCK_DEFAULTS;
+
 import android.os.UserHandle;
 
 import androidx.annotation.NonNull;
 
-import com.android.internal.view.RotationPolicy;
+import com.android.internal.view.RotationPolicy.RotationPolicyListener;
 import com.android.systemui.dagger.SysUISingleton;
+import com.android.systemui.util.wrapper.RotationPolicyWrapper;
 
 import java.util.concurrent.CopyOnWriteArrayList;
 
 import javax.inject.Inject;
+import javax.inject.Named;
 
 /** Platform implementation of the rotation lock controller. **/
 @SysUISingleton
 public final class RotationLockControllerImpl implements RotationLockController {
-    private final Context mContext;
     private final CopyOnWriteArrayList<RotationLockControllerCallback> mCallbacks =
-            new CopyOnWriteArrayList<RotationLockControllerCallback>();
+            new CopyOnWriteArrayList<>();
 
-    private final RotationPolicy.RotationPolicyListener mRotationPolicyListener =
-            new RotationPolicy.RotationPolicyListener() {
+    private final RotationPolicyListener mRotationPolicyListener =
+            new RotationPolicyListener() {
         @Override
         public void onChange() {
             notifyChanged();
         }
     };
 
+    private final RotationPolicyWrapper mRotationPolicy;
+    private final DeviceStateRotationLockSettingController
+            mDeviceStateRotationLockSettingController;
+    private final boolean mIsPerDeviceStateRotationLockEnabled;
+
     @Inject
-    public RotationLockControllerImpl(Context context) {
-        mContext = context;
+    public RotationLockControllerImpl(
+            RotationPolicyWrapper rotationPolicyWrapper,
+            DeviceStateRotationLockSettingController deviceStateRotationLockSettingController,
+            @Named(DEVICE_STATE_ROTATION_LOCK_DEFAULTS) String[] deviceStateRotationLockDefaults
+    ) {
+        mRotationPolicy = rotationPolicyWrapper;
+        mDeviceStateRotationLockSettingController = deviceStateRotationLockSettingController;
+        mIsPerDeviceStateRotationLockEnabled = deviceStateRotationLockDefaults.length > 0;
+        if (mIsPerDeviceStateRotationLockEnabled) {
+            deviceStateRotationLockSettingController.initialize();
+            mCallbacks.add(mDeviceStateRotationLockSettingController);
+        }
+
         setListening(true);
     }
 
@@ -61,32 +79,35 @@ public final class RotationLockControllerImpl implements RotationLockController
     }
 
     public int getRotationLockOrientation() {
-        return RotationPolicy.getRotationLockOrientation(mContext);
+        return mRotationPolicy.getRotationLockOrientation();
     }
 
     public boolean isRotationLocked() {
-        return RotationPolicy.isRotationLocked(mContext);
+        return mRotationPolicy.isRotationLocked();
     }
 
     public void setRotationLocked(boolean locked) {
-        RotationPolicy.setRotationLock(mContext, locked);
+        mRotationPolicy.setRotationLock(locked);
     }
 
     public void setRotationLockedAtAngle(boolean locked, int rotation){
-        RotationPolicy.setRotationLockAtAngle(mContext, locked, rotation);
+        mRotationPolicy.setRotationLockAtAngle(locked, rotation);
     }
 
     public boolean isRotationLockAffordanceVisible() {
-        return RotationPolicy.isRotationLockToggleVisible(mContext);
+        return mRotationPolicy.isRotationLockToggleVisible();
     }
 
     @Override
     public void setListening(boolean listening) {
         if (listening) {
-            RotationPolicy.registerRotationPolicyListener(mContext, mRotationPolicyListener,
+            mRotationPolicy.registerRotationPolicyListener(mRotationPolicyListener,
                     UserHandle.USER_ALL);
         } else {
-            RotationPolicy.unregisterRotationPolicyListener(mContext, mRotationPolicyListener);
+            mRotationPolicy.unregisterRotationPolicyListener(mRotationPolicyListener);
+        }
+        if (mIsPerDeviceStateRotationLockEnabled) {
+            mDeviceStateRotationLockSettingController.setListening(listening);
         }
     }
 
@@ -97,7 +118,7 @@ public final class RotationLockControllerImpl implements RotationLockController
     }
 
     private void notifyChanged(RotationLockControllerCallback callback) {
-        callback.onRotationLockStateChanged(RotationPolicy.isRotationLocked(mContext),
-                RotationPolicy.isRotationLockToggleVisible(mContext));
+        callback.onRotationLockStateChanged(mRotationPolicy.isRotationLocked(),
+                mRotationPolicy.isRotationLockToggleVisible());
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/dagger/StatusBarPolicyModule.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/dagger/StatusBarPolicyModule.java
index c224cf574528c..1eec639670eeb 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/dagger/StatusBarPolicyModule.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/dagger/StatusBarPolicyModule.java
@@ -16,8 +16,10 @@
 
 package com.android.systemui.statusbar.policy.dagger;
 
+import android.content.res.Resources;
 import android.os.UserManager;
 
+import com.android.internal.R;
 import com.android.systemui.dagger.SysUISingleton;
 import com.android.systemui.dagger.qualifiers.Main;
 import com.android.systemui.settings.UserTracker;
@@ -57,6 +59,8 @@ import com.android.systemui.statusbar.policy.ZenModeControllerImpl;
 
 import java.util.concurrent.Executor;
 
+import javax.inject.Named;
+
 import dagger.Binds;
 import dagger.Module;
 import dagger.Provides;
@@ -65,6 +69,9 @@ import dagger.Provides;
 /** Dagger Module for code in the statusbar.policy package. */
 @Module
 public interface StatusBarPolicyModule {
+
+    String DEVICE_STATE_ROTATION_LOCK_DEFAULTS = "DEVICE_STATE_ROTATION_LOCK_DEFAULTS";
+
     /** */
     @Binds
     BluetoothController provideBluetoothController(BluetoothControllerImpl controllerImpl);
@@ -154,4 +161,14 @@ public interface StatusBarPolicyModule {
         controller.init();
         return controller;
     }
+
+    /**
+     * Default values for per-device state rotation lock settings.
+     */
+    @Provides
+    @Named(DEVICE_STATE_ROTATION_LOCK_DEFAULTS)
+    static String[] providesDeviceStateRotationLockDefaults(@Main Resources resources) {
+        return resources.getStringArray(
+                R.array.config_perDeviceStateRotationLockDefaults);
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/util/dagger/UtilModule.java b/packages/SystemUI/src/com/android/systemui/util/dagger/UtilModule.java
index cdfa1457f4a50..981bf01164e3d 100644
--- a/packages/SystemUI/src/com/android/systemui/util/dagger/UtilModule.java
+++ b/packages/SystemUI/src/com/android/systemui/util/dagger/UtilModule.java
@@ -18,12 +18,15 @@ package com.android.systemui.util.dagger;
 
 import com.android.systemui.util.RingerModeTracker;
 import com.android.systemui.util.RingerModeTrackerImpl;
+import com.android.systemui.util.wrapper.UtilWrapperModule;
 
 import dagger.Binds;
 import dagger.Module;
 
 /** Dagger Module for code in the util package. */
-@Module
+@Module(includes = {
+                UtilWrapperModule.class
+        })
 public interface UtilModule {
     /** */
     @Binds
diff --git a/packages/SystemUI/src/com/android/systemui/util/wrapper/RotationPolicyWrapper.kt b/packages/SystemUI/src/com/android/systemui/util/wrapper/RotationPolicyWrapper.kt
new file mode 100644
index 0000000000000..2a0cc7ddacf53
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/util/wrapper/RotationPolicyWrapper.kt
@@ -0,0 +1,67 @@
+/*
+ * 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.
+ */
+
+package com.android.systemui.util.wrapper
+
+import android.content.Context
+import com.android.internal.view.RotationPolicy
+import com.android.internal.view.RotationPolicy.RotationPolicyListener
+import javax.inject.Inject
+
+/**
+ * Testable wrapper interface around RotationPolicy {link com.android.internal.view.RotationPolicy}
+ */
+interface RotationPolicyWrapper {
+    fun setRotationLock(enabled: Boolean)
+    fun setRotationLockAtAngle(enabled: Boolean, rotation: Int)
+    fun getRotationLockOrientation(): Int
+    fun isRotationLockToggleVisible(): Boolean
+    fun isRotationLocked(): Boolean
+    fun registerRotationPolicyListener(listener: RotationPolicyListener, userHandle: Int)
+    fun unregisterRotationPolicyListener(listener: RotationPolicyListener)
+}
+
+class RotationPolicyWrapperImpl @Inject constructor(private val context: Context) :
+    RotationPolicyWrapper {
+
+    override fun setRotationLock(enabled: Boolean) {
+        RotationPolicy.setRotationLock(context, enabled)
+    }
+
+    override fun setRotationLockAtAngle(enabled: Boolean, rotation: Int) {
+        RotationPolicy.setRotationLockAtAngle(context, enabled, rotation)
+    }
+
+    override fun getRotationLockOrientation(): Int =
+        RotationPolicy.getRotationLockOrientation(context)
+
+    override fun isRotationLockToggleVisible(): Boolean =
+        RotationPolicy.isRotationLockToggleVisible(context)
+
+    override fun isRotationLocked(): Boolean =
+        RotationPolicy.isRotationLocked(context)
+
+    override fun registerRotationPolicyListener(
+        listener: RotationPolicyListener,
+        userHandle: Int
+    ) {
+        RotationPolicy.registerRotationPolicyListener(context, listener, userHandle)
+    }
+
+    override fun unregisterRotationPolicyListener(listener: RotationPolicyListener) {
+        RotationPolicy.unregisterRotationPolicyListener(context, listener)
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/util/wrapper/UtilWrapperModule.kt b/packages/SystemUI/src/com/android/systemui/util/wrapper/UtilWrapperModule.kt
new file mode 100644
index 0000000000000..7e3aa277db334
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/util/wrapper/UtilWrapperModule.kt
@@ -0,0 +1,29 @@
+/*
+ * 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.
+ */
+
+package com.android.systemui.util.wrapper
+
+import com.android.systemui.dagger.SysUISingleton
+import dagger.Binds
+import dagger.Module
+
+@Module
+abstract class UtilWrapperModule {
+
+    @Binds
+    @SysUISingleton
+    abstract fun bindRotationPolicyWrapper(impl: RotationPolicyWrapperImpl): RotationPolicyWrapper
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/DeviceStateRotationLockSettingControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/DeviceStateRotationLockSettingControllerTest.java
new file mode 100644
index 0000000000000..30717f431f5b3
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/DeviceStateRotationLockSettingControllerTest.java
@@ -0,0 +1,239 @@
+/*
+ * 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.
+ */
+
+package com.android.systemui.statusbar.policy;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.verify;
+
+import android.hardware.devicestate.DeviceStateManager;
+import android.os.UserHandle;
+import android.provider.Settings;
+import android.testing.AndroidTestingRunner;
+import android.testing.TestableResources;
+
+import androidx.test.filters.SmallTest;
+
+import com.android.internal.view.RotationPolicy;
+import com.android.systemui.SysuiTestCase;
+import com.android.systemui.util.concurrency.FakeExecutor;
+import com.android.systemui.util.settings.FakeSettings;
+import com.android.systemui.util.time.FakeSystemClock;
+import com.android.systemui.util.wrapper.RotationPolicyWrapper;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+@RunWith(AndroidTestingRunner.class)
+@SmallTest
+public class DeviceStateRotationLockSettingControllerTest extends SysuiTestCase {
+
+    private static final String[] DEFAULT_SETTINGS = new String[]{
+            "0:0",
+            "1:2"
+    };
+
+    private final FakeSettings mFakeSettings = new FakeSettings();
+    private final FakeSystemClock mFakeSystemClock = new FakeSystemClock();
+    private final FakeExecutor mFakeExecutor = new FakeExecutor(mFakeSystemClock);
+    @Mock DeviceStateManager mDeviceStateManager;
+    RotationPolicyWrapper mFakeRotationPolicy = new FakeRotationPolicy();
+    DeviceStateRotationLockSettingController mDeviceStateRotationLockSettingController;
+    private DeviceStateManager.DeviceStateCallback mDeviceStateCallback;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(/* testClass= */ this);
+        TestableResources resources = mContext.getOrCreateTestableResources();
+
+        ArgumentCaptor<DeviceStateManager.DeviceStateCallback> deviceStateCallbackArgumentCaptor =
+                ArgumentCaptor.forClass(
+                        DeviceStateManager.DeviceStateCallback.class);
+
+        mDeviceStateRotationLockSettingController = new DeviceStateRotationLockSettingController(
+                mFakeSettings,
+                mFakeRotationPolicy,
+                mDeviceStateManager,
+                mFakeExecutor,
+                DEFAULT_SETTINGS
+        );
+
+        mDeviceStateRotationLockSettingController.setListening(true);
+        verify(mDeviceStateManager).registerCallback(any(),
+                deviceStateCallbackArgumentCaptor.capture());
+        mDeviceStateCallback = deviceStateCallbackArgumentCaptor.getValue();
+    }
+
+    @Test
+    public void whenSavedSettingsEmpty_defaultsLoadedAndSaved() {
+        mFakeSettings.putStringForUser(Settings.Secure.DEVICE_STATE_ROTATION_LOCK, "",
+                UserHandle.USER_CURRENT);
+
+        mDeviceStateRotationLockSettingController.initialize();
+
+        assertThat(mFakeSettings
+                .getStringForUser(Settings.Secure.DEVICE_STATE_ROTATION_LOCK,
+                        UserHandle.USER_CURRENT))
+                .isEqualTo("0:0:1:2");
+    }
+
+    @Test
+    public void whenNoSavedValueForDeviceState_assumeIgnored() {
+        mFakeSettings.putStringForUser(
+                Settings.Secure.DEVICE_STATE_ROTATION_LOCK,
+                /* value= */"0:2:1:2",
+                UserHandle.USER_CURRENT);
+        mFakeRotationPolicy.setRotationLock(true);
+        mDeviceStateRotationLockSettingController.initialize();
+
+        mDeviceStateCallback.onStateChanged(1);
+        assertThat(mFakeRotationPolicy.isRotationLocked()).isFalse();
+
+        // Settings only exist for state 0 and 1
+        mDeviceStateCallback.onStateChanged(2);
+
+        assertThat(mFakeRotationPolicy.isRotationLocked()).isFalse();
+    }
+
+    @Test
+    public void whenDeviceStateSwitched_loadCorrectSetting() {
+        mFakeSettings.putStringForUser(
+                Settings.Secure.DEVICE_STATE_ROTATION_LOCK,
+                /* value= */"0:2:1:1",
+                UserHandle.USER_CURRENT);
+        mFakeRotationPolicy.setRotationLock(true);
+        mDeviceStateRotationLockSettingController.initialize();
+
+        mDeviceStateCallback.onStateChanged(0);
+        assertThat(mFakeRotationPolicy.isRotationLocked()).isFalse();
+
+        mDeviceStateCallback.onStateChanged(1);
+        assertThat(mFakeRotationPolicy.isRotationLocked()).isTrue();
+
+    }
+
+    @Test
+    public void whenUserChangesSetting_saveSettingForCurrentState() {
+        mFakeSettings.putStringForUser(
+                Settings.Secure.DEVICE_STATE_ROTATION_LOCK,
+                /* value= */"0:1:1:2",
+                UserHandle.USER_CURRENT);
+        mFakeRotationPolicy.setRotationLock(true);
+        mDeviceStateRotationLockSettingController.initialize();
+
+        mDeviceStateCallback.onStateChanged(0);
+        assertThat(mFakeRotationPolicy.isRotationLocked()).isTrue();
+
+        mDeviceStateRotationLockSettingController
+                .onRotationLockStateChanged(/* rotationLocked= */false,
+                        /* affordanceVisible= */ true);
+
+        assertThat(mFakeSettings
+                .getStringForUser(Settings.Secure.DEVICE_STATE_ROTATION_LOCK,
+                        UserHandle.USER_CURRENT))
+                .isEqualTo("0:2:1:2");
+    }
+
+
+    @Test
+    public void whenDeviceStateSwitchedToIgnoredState_usePreviousSetting() {
+        mFakeSettings.putStringForUser(
+                Settings.Secure.DEVICE_STATE_ROTATION_LOCK,
+                /* value= */"0:0:1:2",
+                UserHandle.USER_CURRENT);
+        mFakeRotationPolicy.setRotationLock(true);
+        mDeviceStateRotationLockSettingController.initialize();
+
+        mDeviceStateCallback.onStateChanged(1);
+        assertThat(mFakeRotationPolicy.isRotationLocked()).isFalse();
+
+        mDeviceStateCallback.onStateChanged(0);
+        assertThat(mFakeRotationPolicy.isRotationLocked()).isFalse();
+    }
+
+    @Test
+    public void whenDeviceStateSwitchedToIgnoredState_newSettingsSaveForPreviousState() {
+        mFakeSettings.putStringForUser(
+                Settings.Secure.DEVICE_STATE_ROTATION_LOCK,
+                /* value= */"0:0:1:2",
+                UserHandle.USER_CURRENT);
+        mFakeRotationPolicy.setRotationLock(true);
+        mDeviceStateRotationLockSettingController.initialize();
+
+        mDeviceStateCallback.onStateChanged(1);
+        assertThat(mFakeRotationPolicy.isRotationLocked()).isFalse();
+
+        mDeviceStateCallback.onStateChanged(0);
+        assertThat(mFakeRotationPolicy.isRotationLocked()).isFalse();
+
+        mDeviceStateRotationLockSettingController
+                .onRotationLockStateChanged(/* rotationLocked= */true,
+                        /* affordanceVisible= */ true);
+
+        assertThat(mFakeSettings
+                .getStringForUser(Settings.Secure.DEVICE_STATE_ROTATION_LOCK,
+                        UserHandle.USER_CURRENT))
+                .isEqualTo("0:0:1:1");
+    }
+
+    private static class FakeRotationPolicy implements RotationPolicyWrapper {
+
+        private boolean mRotationLock;
+
+        @Override
+        public void setRotationLock(boolean enabled) {
+            mRotationLock = enabled;
+        }
+
+        @Override
+        public void setRotationLockAtAngle(boolean enabled, int rotation) {
+            mRotationLock = enabled;
+        }
+
+        @Override
+        public int getRotationLockOrientation() {
+            throw new AssertionError("Not implemented");
+        }
+
+        @Override
+        public boolean isRotationLockToggleVisible() {
+            throw new AssertionError("Not implemented");
+        }
+
+        @Override
+        public boolean isRotationLocked() {
+            return mRotationLock;
+        }
+
+        @Override
+        public void registerRotationPolicyListener(RotationPolicy.RotationPolicyListener listener,
+                int userHandle) {
+            throw new AssertionError("Not implemented");
+        }
+
+        @Override
+        public void unregisterRotationPolicyListener(
+                RotationPolicy.RotationPolicyListener listener) {
+            throw new AssertionError("Not implemented");
+        }
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/RotationLockControllerImplTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/RotationLockControllerImplTest.java
new file mode 100644
index 0000000000000..0581264d18e2b
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/RotationLockControllerImplTest.java
@@ -0,0 +1,113 @@
+/*
+ * 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.
+ */
+
+package com.android.systemui.statusbar.policy;
+
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyZeroInteractions;
+
+import android.testing.AndroidTestingRunner;
+import android.testing.TestableLooper;
+import android.testing.TestableResources;
+
+import androidx.test.filters.SmallTest;
+
+import com.android.internal.view.RotationPolicy;
+import com.android.systemui.SysuiTestCase;
+import com.android.systemui.util.wrapper.RotationPolicyWrapper;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+@RunWith(AndroidTestingRunner.class)
+@TestableLooper.RunWithLooper
+@SmallTest
+public class RotationLockControllerImplTest extends SysuiTestCase {
+
+    private static final String[] DEFAULT_SETTINGS = new String[]{
+            "0:0",
+            "1:2"
+    };
+
+    @Mock RotationPolicyWrapper mRotationPolicyWrapper;
+    @Mock DeviceStateRotationLockSettingController mDeviceStateRotationLockSettingController;
+
+    private TestableResources mResources;
+    private ArgumentCaptor<RotationPolicy.RotationPolicyListener>
+            mRotationPolicyListenerCaptor;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(/* testClass= */ this);
+        mResources = mContext.getOrCreateTestableResources();
+
+        mRotationPolicyListenerCaptor = ArgumentCaptor.forClass(
+                RotationPolicy.RotationPolicyListener.class);
+    }
+
+    @Test
+    public void whenFlagOff_doesntInteractWithDeviceStateRotationController() {
+        createRotationLockController(new String[0]);
+
+        verifyZeroInteractions(mDeviceStateRotationLockSettingController);
+    }
+
+    @Test
+    public void whenFlagOn_setListeningSetsListeningOnDeviceStateRotationController() {
+        createRotationLockController();
+
+        verify(mDeviceStateRotationLockSettingController).setListening(/* listening= */ true);
+    }
+
+    @Test
+    public void whenFlagOn_initializesDeviceStateRotationController() {
+        createRotationLockController();
+
+        verify(mDeviceStateRotationLockSettingController).initialize();
+    }
+
+    @Test
+    public void whenFlagOn_dviceStateRotationControllerAddedToCallbacks() {
+        createRotationLockController();
+        captureRotationPolicyListener().onChange();
+
+        verify(mDeviceStateRotationLockSettingController)
+                .onRotationLockStateChanged(anyBoolean(), anyBoolean());
+    }
+
+    private RotationPolicy.RotationPolicyListener captureRotationPolicyListener() {
+        verify(mRotationPolicyWrapper)
+                .registerRotationPolicyListener(mRotationPolicyListenerCaptor.capture(), anyInt());
+        return mRotationPolicyListenerCaptor.getValue();
+    }
+
+    private void createRotationLockController() {
+        createRotationLockController(DEFAULT_SETTINGS);
+    }
+    private void createRotationLockController(String[] deviceStateRotationLockDefaults) {
+        new RotationLockControllerImpl(
+                mRotationPolicyWrapper,
+                mDeviceStateRotationLockSettingController,
+                deviceStateRotationLockDefaults
+        );
+    }
+}