[pvmfw] Move dirty state management control util to vmbase

As part of the project of moving the memory management code from
pvmfw to vmbase so that we can reuse it in rialto.

Bug: 284462758
Test: m pvmfw_img
Change-Id: I2cf71f04551af500627c2777b37d7e5a7d8b4a1d
diff --git a/vmbase/Android.bp b/vmbase/Android.bp
index 5ed436c..237a377 100644
--- a/vmbase/Android.bp
+++ b/vmbase/Android.bp
@@ -70,6 +70,9 @@
     whole_static_libs: [
         "librust_baremetal",
     ],
+    features: [
+        "cpu_feat_hafdbs",
+    ],
     apex_available: ["com.android.virt"],
 }
 
diff --git a/vmbase/src/lib.rs b/vmbase/src/lib.rs
index 0eca1e4..d086f1c 100644
--- a/vmbase/src/lib.rs
+++ b/vmbase/src/lib.rs
@@ -23,6 +23,7 @@
 pub mod layout;
 mod linker;
 pub mod logger;
+pub mod memory;
 pub mod power;
 pub mod uart;
 
diff --git a/vmbase/src/memory/dbm.rs b/vmbase/src/memory/dbm.rs
new file mode 100644
index 0000000..3a52c2d
--- /dev/null
+++ b/vmbase/src/memory/dbm.rs
@@ -0,0 +1,47 @@
+// Copyright 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.
+
+//! Hardware management of the access flag and dirty state.
+
+use crate::{isb, read_sysreg, write_sysreg};
+
+/// Sets whether the hardware management of access and dirty state is enabled with
+/// the given boolean.
+pub fn set_dbm_enabled(enabled: bool) {
+    if !dbm_available() {
+        return;
+    }
+    // TCR_EL1.{HA,HD} bits controlling hardware management of access and dirty state
+    const TCR_EL1_HA_HD_BITS: usize = 3 << 39;
+
+    let mut tcr = read_sysreg!("tcr_el1");
+    if enabled {
+        tcr |= TCR_EL1_HA_HD_BITS
+    } else {
+        tcr &= !TCR_EL1_HA_HD_BITS
+    };
+    // Safe because it writes to a system register and does not affect Rust.
+    unsafe { write_sysreg!("tcr_el1", tcr) }
+    isb!();
+}
+
+/// Returns `true` if hardware dirty state management is available.
+fn dbm_available() -> bool {
+    if !cfg!(feature = "cpu_feat_hafdbs") {
+        return false;
+    }
+    // Hardware dirty bit management available flag (ID_AA64MMFR1_EL1.HAFDBS[1])
+    const DBM_AVAILABLE: usize = 1 << 1;
+    read_sysreg!("id_aa64mmfr1_el1") & DBM_AVAILABLE != 0
+}
diff --git a/vmbase/src/memory/mod.rs b/vmbase/src/memory/mod.rs
new file mode 100644
index 0000000..c4990f2
--- /dev/null
+++ b/vmbase/src/memory/mod.rs
@@ -0,0 +1,19 @@
+// Copyright 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.
+
+//! Memory management.
+
+mod dbm;
+
+pub use dbm::set_dbm_enabled;