refactor(lib/realm): optimize s2tt component

This patch introduces optimizations into the s2tt and rtt components
in order to improve performance. It also fixes most of the MISRA
warnings already existing on the modified files.

Signed-off-by: Javier Almansa Sobrino <javier.almansasobrino@arm.com>
Change-Id: I8c96e76badeca560db316853ab1730df332a3474
diff --git a/lib/realm/src/s2tt.c b/lib/realm/src/s2tt.c
index 49fe430..130ffb1 100644
--- a/lib/realm/src/s2tt.c
+++ b/lib/realm/src/s2tt.c
@@ -11,121 +11,15 @@
 #include <memory_alloc.h>
 #include <realm.h>
 #include <ripas.h>
+#include <s2tt.h>
+#include <s2tt_pvt_defs.h>
 #include <smc.h>
 #include <status.h>
 #include <stddef.h>
 #include <string.h>
-#include <table.h>
 
-/*
- * For prototyping we assume 4K pages
- */
-#define BLOCK_L2_SIZE		(GRANULE_SIZE * S2TTES_PER_S2TT)
-
-/*
- * The maximum number of bits supported by the RMM for a stage 2 translation
- * output address (including stage 2 table entries).
- */
-#define S2TTE_OA_BITS			48U
-
-#define DESC_TYPE_MASK			3UL
-#define S2TTE_Lx_INVALID		0UL
-#define S2TTE_L012_BLOCK		1UL
-#define S2TTE_L012_TABLE		3UL
-#define S2TTE_L3_PAGE			3UL
-
-/*
- * The following constants for the mapping attributes (S2_TTE_MEMATTR_*)
- * assume that HCR_EL2.FWB is set.
- */
-#define S2TTE_MEMATTR_SHIFT		2
-#define S2TTE_MEMATTR_MASK		(0x7UL << S2TTE_MEMATTR_SHIFT)
-#define S2TTE_MEMATTR_FWB_NORMAL_WB	((1UL << 4) | (2UL << 2))
-#define S2TTE_MEMATTR_FWB_RESERVED	((1UL << 4) | (0UL << 2))
-
-#define S2TTE_AP_SHIFT			6
-#define S2TTE_AP_MASK			(3UL << S2TTE_AP_SHIFT)
-#define S2TTE_AP_RW			(3UL << S2TTE_AP_SHIFT)
-
-#define S2TTE_SH_SHIFT			8
-#define S2TTE_SH_MASK			(3UL << S2TTE_SH_SHIFT)
-#define S2TTE_SH_NS			(0UL << S2TTE_SH_SHIFT)
-#define S2TTE_SH_RESERVED		(1UL << S2TTE_SH_SHIFT)
-#define S2TTE_SH_OS			(2UL << S2TTE_SH_SHIFT)
-#define S2TTE_SH_IS			(3UL << S2TTE_SH_SHIFT)	/* Inner Shareable */
-
-/*
- * We set HCR_EL2.FWB So we set bit[4] to 1 and bits[3:2] to 2 and force
- * everyting to be Normal Write-Back
- */
-#define S2TTE_MEMATTR_FWB_NORMAL_WB	((1UL << 4) | (2UL << 2))
-#define S2TTE_AF			(1UL << 10)
-#define S2TTE_XN			(2UL << 53)
-#define S2TTE_NS			(1UL << 55)
-
-#define S2TTE_ATTRS	(S2TTE_MEMATTR_FWB_NORMAL_WB | S2TTE_AP_RW | \
-			S2TTE_SH_IS | S2TTE_AF)
-#define S2TTE_NS_ATTR_MASK (S2TTE_MEMATTR_MASK | S2TTE_AP_MASK | \
-			    S2TTE_SH_MASK)
-
-#define S2TTE_TABLE	S2TTE_L012_TABLE
-#define S2TTE_BLOCK	(S2TTE_ATTRS | S2TTE_L012_BLOCK)
-#define S2TTE_PAGE	(S2TTE_ATTRS | S2TTE_L3_PAGE)
-#define S2TTE_BLOCK_NS	(S2TTE_NS | S2TTE_XN | S2TTE_AF | S2TTE_L012_BLOCK)
-#define S2TTE_PAGE_NS	(S2TTE_NS | S2TTE_XN | S2TTE_AF | S2TTE_L3_PAGE)
-#define S2TTE_INVALID	S2TTE_Lx_INVALID
-
-/*
- * The type of stage 2 translation table entry (s2tte) is defined by:
- * 1. Table level where it resides
- * 2. DESC_TYPE field[1:0]
- * 4. HIPAS field [4:2]
- * 4. RIPAS field [6:5]
- * 5. NS field [55]
- *
- * ======================================================================================
- * s2tte type           level DESC_TYPE[1:0] HIPAS[4:2]    RIPAS[6:5]   NS  OA alignment
- * ======================================================================================
- * unassigned_empty     any   invalid[0]     unassigned[0] empty[0]     0   n/a
- * --------------------------------------------------------------------------------------
- * unassigned_ram       any   invalid[0]     unassigned[0] ram[1]       0   n/a
- * --------------------------------------------------------------------------------------
- * unassigned_destroyed any   invalid[0]     unassigned[0] destroyed[2] 0   n/a
- * --------------------------------------------------------------------------------------
- * assigned_empty       2,3   invalid[0]     assigned[1]   empty[0]     0   to level
- * --------------------------------------------------------------------------------------
- * assigned_ram         3     page[3]        n/a           n/a          0   to level
- *                      2     block[1]       n/a           n/a          0   to level
- * --------------------------------------------------------------------------------------
- * assigned_destroyed   any   invalid[0]     assigned[1]   destroyed[2] 0   n/a
- * ======================================================================================
- * unassigned_ns        any   invalid[0]     unassigned[0] n/a          1   n/a
- * --------------------------------------------------------------------------------------
- * assigned_ns	        3     page[3]        n/a           n/a          1   to level
- *                      2     block[1]       n/a           n/a          1   to level
- * ======================================================================================
- * table              <=2     table[3]       n/a           n/a          n/a to 4K
- * ======================================================================================
- */
-
-#define S2TTE_INVALID_HIPAS_SHIFT	2
-#define S2TTE_INVALID_HIPAS_WIDTH	3U
-#define S2TTE_INVALID_HIPAS_MASK	MASK(S2TTE_INVALID_HIPAS)
-
-#define S2TTE_INVALID_HIPAS_UNASSIGNED	(INPLACE(S2TTE_INVALID_HIPAS, 0UL))
-#define S2TTE_INVALID_HIPAS_ASSIGNED	(INPLACE(S2TTE_INVALID_HIPAS, 1UL))
-
-#define S2TTE_INVALID_RIPAS_SHIFT	5
-#define S2TTE_INVALID_RIPAS_WIDTH	2U
-#define S2TTE_INVALID_RIPAS_MASK	MASK(S2TTE_INVALID_RIPAS)
-
-#define S2TTE_INVALID_RIPAS_EMPTY	(INPLACE(S2TTE_INVALID_RIPAS, 0UL))
-#define S2TTE_INVALID_RIPAS_RAM		(INPLACE(S2TTE_INVALID_RIPAS, 1UL))
-#define S2TTE_INVALID_RIPAS_DESTROYED	(INPLACE(S2TTE_INVALID_RIPAS, 2UL))
-
-#define S2TTE_INVALID_UNPROTECTED	0x0UL
-
-#define NR_RTT_LEVELS	4
+/* TODO: Fix this when introducing LPA2 support */
+COMPILER_ASSERT(S2TT_MIN_STARTING_LEVEL >= 0);
 
 /*
  * Invalidates S2 TLB entries from [ipa, ipa + size] region tagged with `vmid`.
@@ -147,6 +41,8 @@
 	 *   - Address range invalidation.
 	 */
 
+	assert(s2_ctx != NULL);
+
 	/*
 	 * Save the current content of vttb_el2.
 	 */
@@ -191,7 +87,7 @@
  * Call this function after:
  * 1.  A L3 page desc has been removed.
  */
-void invalidate_page(const struct realm_s2_context *s2_ctx, unsigned long addr)
+void s2tt_invalidate_page(const struct realm_s2_context *s2_ctx, unsigned long addr)
 {
 	stage2_tlbi_ipa(s2_ctx, addr, GRANULE_SIZE);
 }
@@ -203,7 +99,7 @@
  * 2a. A L2 table desc has been removed, where
  * 2b. All S2TTEs in L3 table that the L2 table desc was pointed to were invalid.
  */
-void invalidate_block(const struct realm_s2_context *s2_ctx, unsigned long addr)
+void s2tt_invalidate_block(const struct realm_s2_context *s2_ctx, unsigned long addr)
 {
 	stage2_tlbi_ipa(s2_ctx, addr, GRANULE_SIZE);
 }
@@ -214,7 +110,7 @@
  * 1a. A L2 table desc has been removed, where
  * 1b. Some S2TTEs in the table that the L2 table desc was pointed to were valid.
  */
-void invalidate_pages_in_block(const struct realm_s2_context *s2_ctx, unsigned long addr)
+void s2tt_invalidate_pages_in_block(const struct realm_s2_context *s2_ctx, unsigned long addr)
 {
 	stage2_tlbi_ipa(s2_ctx, addr, BLOCK_L2_SIZE);
 }
@@ -232,9 +128,7 @@
 {
 	unsigned int levels, lsb;
 
-	assert(level <= RTT_PAGE_LEVEL);
-
-	levels = (unsigned int)(RTT_PAGE_LEVEL - level);
+	levels = (unsigned int)(S2TT_PAGE_LEVEL - level);
 	lsb = (levels * S2TTE_STRIDE) + GRANULE_SHIFT;
 
 	addr >>= lsb;
@@ -257,9 +151,7 @@
 {
 	unsigned int levels, lsb;
 
-	assert(start_level <= RTT_PAGE_LEVEL);
-
-	levels = (unsigned int)(RTT_PAGE_LEVEL - start_level);
+	levels = (unsigned int)(S2TT_PAGE_LEVEL - start_level);
 	lsb = (levels * S2TTE_STRIDE) + GRANULE_SHIFT;
 
 	addr &= (1UL << ipa_bits) - 1UL;
@@ -271,27 +163,23 @@
 {
 	unsigned int levels, lsb, msb;
 
-	assert(level <= RTT_PAGE_LEVEL);
+	assert(level <= S2TT_PAGE_LEVEL);
+	assert(level >= S2TT_MIN_STARTING_LEVEL);
 
-	levels = (unsigned int)(RTT_PAGE_LEVEL - level);
+	levels = (unsigned int)(S2TT_PAGE_LEVEL - level);
 	lsb = (levels * S2TTE_STRIDE) + GRANULE_SHIFT;
 	msb = S2TTE_OA_BITS - 1U;
 
 	return (addr & BIT_MASK_ULL(msb, lsb));
 }
 
-static inline unsigned long table_entry_to_phys(unsigned long entry)
-{
-	return addr_level_mask(entry, RTT_PAGE_LEVEL);
-}
-
 static inline bool entry_is_table(unsigned long entry)
 {
-	return ((entry & DESC_TYPE_MASK) == S2TTE_L012_TABLE);
+	return ((entry & S2TT_DESC_TYPE_MASK) == S2TTE_L012_TABLE);
 }
 
-static unsigned long __table_get_entry(struct granule *g_tbl,
-				       unsigned long idx)
+static unsigned long table_get_entry(struct granule *g_tbl,
+				     unsigned long idx)
 {
 	unsigned long *table, entry;
 
@@ -304,10 +192,12 @@
 	return entry;
 }
 
-static struct granule *__find_next_level_idx(struct granule *g_tbl,
-					     unsigned long idx)
+#define table_entry_to_phys(tte)	addr_level_mask(tte, S2TT_PAGE_LEVEL)
+
+static struct granule *find_next_level_idx(struct granule *g_tbl,
+					   unsigned long idx)
 {
-	const unsigned long entry = __table_get_entry(g_tbl, idx);
+	const unsigned long entry = table_get_entry(g_tbl, idx);
 
 	if (!entry_is_table(entry)) {
 		return NULL;
@@ -316,12 +206,12 @@
 	return addr_to_granule(table_entry_to_phys(entry));
 }
 
-static struct granule *__find_lock_next_level(struct granule *g_tbl,
-					      unsigned long map_addr,
-					      long level)
+static struct granule *find_lock_next_level(struct granule *g_tbl,
+					    unsigned long map_addr,
+					    long level)
 {
 	const unsigned long idx = s2_addr_to_idx(map_addr, level);
-	struct granule *g = __find_next_level_idx(g_tbl, idx);
+	struct granule *g = find_next_level_idx(g_tbl, idx);
 
 	if (g != NULL) {
 		granule_lock(g, GRANULE_STATE_RTT);
@@ -346,32 +236,40 @@
  * - Level @level is reached.
  *
  * On return:
- * - rtt_walk::last_level is the last level that has been reached by the walk.
- * - rtt_walk.g_llt points to the TABLE granule at level @rtt_walk::level.
+ * - s2tt_walk::last_level is the last level that has been reached by the walk.
+ * - s2tt_walk.g_llt points to the TABLE granule at level @s2tt_walk::level.
  *   The granule is locked.
- * - rtt_walk::index is the entry index at rtt_walk.g_llt for @map_addr.
+ * - s2tt_walk::index is the entry index at s2tt_walk.g_llt for @map_addr.
  */
-void rtt_walk_lock_unlock(struct granule *g_root,
+void s2tt_walk_lock_unlock(struct granule *g_root,
 			  int start_level,
 			  unsigned long ipa_bits,
 			  unsigned long map_addr,
 			  long level,
-			  struct rtt_walk *wi)
+			  struct s2tt_walk *wi)
 {
-	struct granule *g_tbls[NR_RTT_LEVELS] = { NULL };
+	struct granule *g_tbls[NR_RTT_LEVELS] = { (struct granule *)NULL };
 	unsigned long sl_idx;
 	int i, last_level;
 
-	assert(start_level >= MIN_STARTING_LEVEL);
+	assert(start_level >= S2TT_MIN_STARTING_LEVEL);
+	assert(start_level <= S2TT_PAGE_LEVEL);
 	assert(level >= start_level);
+	assert(level <= S2TT_PAGE_LEVEL);
 	assert(map_addr < (1UL << ipa_bits));
 	assert(wi != NULL);
+	assert(g_root != NULL);
+	assert(ipa_bits <= s2tt_max_ipa_size());
 
 	/* Handle concatenated starting level (SL) tables */
 	sl_idx = s2_sl_addr_to_idx(map_addr, start_level, ipa_bits);
 	if (sl_idx >= S2TTES_PER_S2TT) {
 		unsigned int tt_num = (unsigned int)(sl_idx >> S2TTE_STRIDE);
-		struct granule *g_concat_root = (struct granule *)((uintptr_t)g_root +
+		struct granule *g_concat_root;
+
+		assert(tt_num < S2TTE_MAX_CONCAT_TABLES);
+
+		g_concat_root = (struct granule *)((uintptr_t)g_root +
 						(tt_num * sizeof(struct granule)));
 
 		granule_lock(g_concat_root, GRANULE_STATE_RTT);
@@ -387,7 +285,7 @@
 		 * (previous level). Also, hand-over-hand locking/unlocking is
 		 * used to avoid race conditions.
 		 */
-		g_tbls[i + 1] = __find_lock_next_level(g_tbls[i], map_addr, i);
+		g_tbls[i + 1] = find_lock_next_level(g_tbls[i], map_addr, i);
 		if (g_tbls[i + 1] == NULL) {
 			last_level = i;
 			goto out;
@@ -427,29 +325,54 @@
 }
 
 /*
- * Creates an invalid s2tte with output address @pa, HIPAS=ASSIGNED and
- * RIPAS=DESTROYED, at level @level.
+ * Creates an unassigned_ns s2tte.
  */
-unsigned long s2tte_create_assigned_destroyed(unsigned long pa, long level)
+unsigned long s2tte_create_unassigned_ns(void)
 {
-	(void)level;
-
-	assert(level >= RTT_MIN_BLOCK_LEVEL);
-	assert(addr_is_level_aligned(pa, level));
-	return (pa | S2TTE_INVALID_HIPAS_ASSIGNED | S2TTE_INVALID_RIPAS_DESTROYED);
+	return (S2TTE_NS | S2TTE_INVALID_HIPAS_UNASSIGNED |
+		S2TTE_INVALID_UNPROTECTED);
 }
 
 /*
  * Creates an invalid s2tte with output address @pa, HIPAS=ASSIGNED and
- * RIPAS=EMPTY, at level @level.
+ * RIPAS=@s2tte_ripas, at level @level.
  */
-unsigned long s2tte_create_assigned_empty(unsigned long pa, long level)
+static unsigned long s2tte_create_assigned(unsigned long pa, long level,
+					   unsigned long s2tte_ripas)
 {
 	(void)level;
 
-	assert(level >= RTT_MIN_BLOCK_LEVEL);
-	assert(addr_is_level_aligned(pa, level));
-	return (pa | S2TTE_INVALID_HIPAS_ASSIGNED | S2TTE_INVALID_RIPAS_EMPTY);
+	assert(level >= S2TT_MIN_BLOCK_LEVEL);
+	assert(level <= S2TT_PAGE_LEVEL);
+	assert(s2tte_ripas <= S2TTE_INVALID_RIPAS_DESTROYED);
+	assert(s2tte_is_addr_lvl_aligned(pa, level));
+
+	if (s2tte_ripas == S2TTE_INVALID_RIPAS_RAM) {
+		if (level == S2TT_PAGE_LEVEL) {
+			return (pa | S2TTE_PAGE);
+		}
+		return (pa | S2TTE_BLOCK);
+	}
+
+	return (pa | S2TTE_INVALID_HIPAS_ASSIGNED | s2tte_ripas);
+}
+
+/*
+ * Creates and invalid s2tte with output address @pa, HIPAS=ASSIGNED and
+ * RIPAS=DESTROYED at level @level.
+ */
+unsigned long s2tte_create_assigned_destroyed(unsigned long pa, long level)
+{
+	return s2tte_create_assigned(pa, level, S2TTE_INVALID_RIPAS_DESTROYED);
+}
+
+/*
+ * Creates an invalid s2tte with output address @pa, HIPAS=ASSIGNED and
+ * RIPAS=EMPTY at level @level.
+ */
+unsigned long s2tte_create_assigned_empty(unsigned long pa, long level)
+{
+	return s2tte_create_assigned(pa, level, S2TTE_INVALID_RIPAS_EMPTY);
 }
 
 /*
@@ -457,12 +380,7 @@
  */
 unsigned long s2tte_create_assigned_ram(unsigned long pa, long level)
 {
-	assert(level >= RTT_MIN_BLOCK_LEVEL);
-	assert(addr_is_level_aligned(pa, level));
-	if (level == RTT_PAGE_LEVEL) {
-		return (pa | S2TTE_PAGE);
-	}
-	return (pa | S2TTE_BLOCK);
+	return s2tte_create_assigned(pa, level, S2TTE_INVALID_RIPAS_RAM);
 }
 
 /*
@@ -475,25 +393,7 @@
 {
 	unsigned long current_ripas = s2tte & S2TTE_INVALID_RIPAS_MASK;
 
-	assert(level >= RTT_MIN_BLOCK_LEVEL);
-	assert(addr_is_level_aligned(pa, level));
-	assert(EXTRACT(S2TTE_INVALID_RIPAS, current_ripas) <=
-	       EXTRACT(S2TTE_INVALID_RIPAS, S2TTE_INVALID_RIPAS_DESTROYED));
-
-	if (current_ripas != S2TTE_INVALID_RIPAS_RAM) {
-		return (pa | S2TTE_INVALID_HIPAS_ASSIGNED | current_ripas);
-	}
-
-	return s2tte_create_assigned_ram(pa, level);
-}
-
-/*
- * Creates an unassigned_ns s2tte.
- */
-unsigned long s2tte_create_unassigned_ns(void)
-{
-	return (S2TTE_NS | S2TTE_INVALID_HIPAS_UNASSIGNED |
-		S2TTE_INVALID_UNPROTECTED);
+	return s2tte_create_assigned(pa, level, current_ripas);
 }
 
 /*
@@ -504,14 +404,16 @@
  * - MemAttr
  * - S2AP
  * - Shareability
- * Any other field on @s2tte is masked out.
  */
 unsigned long s2tte_create_assigned_ns(unsigned long s2tte, long level)
 {
-	unsigned long new_s2tte = s2tte & ~DESC_TYPE_MASK;
+	unsigned long new_s2tte = s2tte & ~S2TT_DESC_TYPE_MASK;
 
-	assert(level >= RTT_MIN_BLOCK_LEVEL);
-	if (level == RTT_PAGE_LEVEL) {
+	assert(level >= S2TT_MIN_BLOCK_LEVEL);
+	assert(level <= S2TT_PAGE_LEVEL);
+	assert((s2tte & S2TTE_NS_ATTR_RMM) == 0UL);
+
+	if (level == S2TT_PAGE_LEVEL) {
 		return (new_s2tte | S2TTE_PAGE_NS);
 	}
 	return (new_s2tte | S2TTE_BLOCK_NS);
@@ -522,8 +424,11 @@
  */
 bool host_ns_s2tte_is_valid(unsigned long s2tte, long level)
 {
+
 	unsigned long mask = addr_level_mask(~0UL, level) |
-						S2TTE_NS_ATTR_MASK;
+						S2TTE_NS_ATTR_HOST_MASK;
+
+	assert(level >= S2TT_MIN_BLOCK_LEVEL);
 
 	/*
 	 * Test that all fields that are not controlled by the host are zero
@@ -544,7 +449,7 @@
 	/*
 	 * Only one value masked by S2TTE_SH_MASK is invalid/reserved.
 	 */
-	if ((s2tte & S2TTE_SH_MASK) == S2TTE_SH_RESERVED) {
+	if ((s2tte & S2TTE_SH_MASK) != S2TTE_SH_IS) {
 		return false;
 	}
 
@@ -560,7 +465,9 @@
 unsigned long host_ns_s2tte(unsigned long s2tte, long level)
 {
 	unsigned long mask = addr_level_mask(~0UL, level) |
-						S2TTE_NS_ATTR_MASK;
+						S2TTE_NS_ATTR_HOST_MASK;
+	assert(level >= S2TT_MIN_BLOCK_LEVEL);
+
 	return (s2tte & mask);
 }
 
@@ -571,7 +478,8 @@
 {
 	(void)level;
 
-	assert(level < RTT_PAGE_LEVEL);
+	assert(level < S2TT_PAGE_LEVEL);
+	assert(level >= S2TT_MIN_STARTING_LEVEL);
 	assert(GRANULE_ALIGNED(pa));
 
 	return (pa | S2TTE_TABLE);
@@ -594,9 +502,9 @@
 /*
  * Returns true if @s2tte has HIPAS=@hipas.
  */
-static bool s2tte_has_hipas(unsigned long s2tte, unsigned long hipas)
+static inline bool s2tte_has_hipas(unsigned long s2tte, unsigned long hipas)
 {
-	unsigned long desc_type = s2tte & DESC_TYPE_MASK;
+	unsigned long desc_type = s2tte & S2TT_DESC_TYPE_MASK;
 	unsigned long invalid_desc_hipas = s2tte & S2TTE_INVALID_HIPAS_MASK;
 
 	return ((desc_type == S2TTE_INVALID) && (invalid_desc_hipas == hipas));
@@ -605,7 +513,8 @@
 /*
  * Returns true if @s2tte has HIPAS=UNASSIGNED and RIPAS=@ripas.
  */
-static bool s2tte_has_unassigned_ripas(unsigned long s2tte, unsigned long ripas)
+static inline bool s2tte_has_unassigned_ripas(unsigned long s2tte,
+					      unsigned long ripas)
 {
 	return (((s2tte & S2TTE_INVALID_RIPAS_MASK) == ripas) &&
 		  s2tte_has_hipas(s2tte, S2TTE_INVALID_HIPAS_UNASSIGNED));
@@ -614,7 +523,8 @@
 /*
  * Returns true if @s2tte has HIPAS=ASSIGNED and RIPAS=@ripas.
  */
-static bool s2tte_has_assigned_ripas(unsigned long s2tte, unsigned long ripas)
+static inline bool s2tte_has_assigned_ripas(unsigned long s2tte,
+					    unsigned long ripas)
 {
 	return (((s2tte & S2TTE_INVALID_RIPAS_MASK) == ripas) &&
 		  s2tte_has_hipas(s2tte, S2TTE_INVALID_HIPAS_ASSIGNED));
@@ -690,11 +600,11 @@
 		return false;
 	}
 
-	desc_type = s2tte & DESC_TYPE_MASK;
+	desc_type = s2tte & S2TT_DESC_TYPE_MASK;
 
 	/* Only pages at L3 and valid blocks at L2 allowed */
-	if (((level == RTT_PAGE_LEVEL) && (desc_type == S2TTE_L3_PAGE)) ||
-	    ((level == RTT_MIN_BLOCK_LEVEL) && (desc_type == S2TTE_L012_BLOCK))) {
+	if (((level == S2TT_PAGE_LEVEL) && (desc_type == S2TTE_L3_PAGE)) ||
+	    ((level == S2TT_MIN_BLOCK_LEVEL) && (desc_type == S2TTE_L012_BLOCK))) {
 		return true;
 	}
 
@@ -722,15 +632,15 @@
  */
 bool s2tte_is_table(unsigned long s2tte, long level)
 {
-	return ((level < RTT_PAGE_LEVEL) &&
-		((s2tte & DESC_TYPE_MASK) == S2TTE_L012_TABLE));
+	return ((level < S2TT_PAGE_LEVEL) &&
+		((s2tte & S2TT_DESC_TYPE_MASK) == S2TTE_L012_TABLE));
 }
 
 /*
  * Returns RIPAS of @s2tte.
  *
  * Caller should ensure that HIPAS=UNASSIGNED or HIPAS=ASSIGNED.
- * The s2tte must be not valid/invalid descriptor.
+ * The s2tte, if valid, should correspond to RIPAS_RAM.
  */
 enum ripas s2tte_get_ripas(unsigned long s2tte)
 {
@@ -741,20 +651,17 @@
 	 * bit is 1 (S2AP is set to RW for lower EL), which corresponds
 	 * to RIPAS_RAM (bits[6:5] = b01) on a valid descriptor.
 	 */
-	if (((s2tte & DESC_TYPE_MASK) != S2TTE_INVALID) &&
-	     (desc_ripas != S2TTE_INVALID_RIPAS_RAM)) {
-		assert(false);
-	}
+	assert(((s2tte & S2TT_DESC_TYPE_MASK) == S2TTE_INVALID) ||
+	       (desc_ripas == S2TTE_INVALID_RIPAS_RAM));
 
-	switch (desc_ripas) {
-	case S2TTE_INVALID_RIPAS_EMPTY:
-		return RIPAS_EMPTY;
-	case S2TTE_INVALID_RIPAS_RAM:
-		return RIPAS_RAM;
-	default:
-		assert(desc_ripas == S2TTE_INVALID_RIPAS_DESTROYED);
-		return RIPAS_DESTROYED;
-	}
+	assert(EXTRACT(S2TTE_INVALID_HIPAS, s2tte) <=
+	       EXTRACT(S2TTE_INVALID_HIPAS, S2TTE_INVALID_HIPAS_ASSIGNED));
+
+	desc_ripas = desc_ripas >> S2TTE_INVALID_RIPAS_SHIFT;
+
+	assert(desc_ripas <= (unsigned int)RIPAS_DESTROYED);
+
+	return (enum ripas)desc_ripas;
 }
 
 /*
@@ -765,8 +672,13 @@
  */
 void s2tt_init_unassigned_empty(unsigned long *s2tt)
 {
+	assert(s2tt != NULL);
+
+	unsigned long s2tte =
+		s2tte_create_unassigned_empty();
+
 	for (unsigned int i = 0U; i < S2TTES_PER_S2TT; i++) {
-		s2tt[i] = s2tte_create_unassigned_empty();
+		s2tt[i] = s2tte;
 	}
 
 	dsb(ish);
@@ -780,8 +692,12 @@
  */
 void s2tt_init_unassigned_ram(unsigned long *s2tt)
 {
+	assert(s2tt != NULL);
+
+	unsigned long s2tte = s2tte_create_unassigned_ram();
+
 	for (unsigned int i = 0U; i < S2TTES_PER_S2TT; i++) {
-		s2tt[i] = s2tte_create_unassigned_ram();
+		s2tt[i] = s2tte;
 	}
 
 	dsb(ish);
@@ -795,8 +711,12 @@
  */
 void s2tt_init_unassigned_ns(unsigned long *s2tt)
 {
+	assert(s2tt != NULL);
+
+	unsigned long s2tte = s2tte_create_unassigned_ns();
+
 	for (unsigned int i = 0U; i < S2TTES_PER_S2TT; i++) {
-		s2tt[i] = s2tte_create_unassigned_ns();
+		s2tt[i] = s2tte;
 	}
 
 	dsb(ish);
@@ -810,8 +730,13 @@
  */
 void s2tt_init_unassigned_destroyed(unsigned long *s2tt)
 {
+	assert(s2tt != NULL);
+
+	unsigned long s2tte =
+		s2tte_create_unassigned_destroyed();
+
 	for (unsigned int i = 0U; i < S2TTES_PER_S2TT; i++) {
-		s2tt[i] = s2tte_create_unassigned_destroyed();
+		s2tt[i] = s2tte;
 	}
 	dsb(ish);
 }
@@ -825,6 +750,11 @@
  */
 void s2tt_init_assigned_destroyed(unsigned long *s2tt, unsigned long pa, long level)
 {
+	assert(level >= S2TT_MIN_BLOCK_LEVEL);
+	assert(level <= S2TT_PAGE_LEVEL);
+	assert(s2tt != NULL);
+	assert(s2tte_is_addr_lvl_aligned(pa, level));
+
 	const unsigned long map_size = s2tte_map_size(level);
 
 	for (unsigned int i = 0U; i < S2TTES_PER_S2TT; i++) {
@@ -834,17 +764,6 @@
 	dsb(ish);
 }
 
-unsigned long s2tte_map_size(long level)
-{
-	unsigned int levels, lsb;
-
-	assert(level <= RTT_PAGE_LEVEL);
-
-	levels = (unsigned int)(RTT_PAGE_LEVEL - level);
-	lsb = (levels * S2TTE_STRIDE) + GRANULE_SHIFT;
-	return (1UL << lsb);
-}
-
 /*
  * Populates @s2tt with HIPAS=ASSIGNED, RIPAS=EMPTY s2ttes that refer to a
  * contiguous memory block starting at @pa, and mapped at level @level.
@@ -854,6 +773,11 @@
  */
 void s2tt_init_assigned_empty(unsigned long *s2tt, unsigned long pa, long level)
 {
+	assert(level >= S2TT_MIN_BLOCK_LEVEL);
+	assert(level <= S2TT_PAGE_LEVEL);
+	assert(s2tt != NULL);
+	assert(s2tte_is_addr_lvl_aligned(pa, level));
+
 	const unsigned long map_size = s2tte_map_size(level);
 
 	for (unsigned int i = 0U; i < S2TTES_PER_S2TT; i++) {
@@ -872,6 +796,11 @@
  */
 void s2tt_init_assigned_ram(unsigned long *s2tt, unsigned long pa, long level)
 {
+	assert(level >= S2TT_MIN_BLOCK_LEVEL);
+	assert(level <= S2TT_PAGE_LEVEL);
+	assert(s2tt != NULL);
+	assert(s2tte_is_addr_lvl_aligned(pa, level));
+
 	const unsigned long map_size = s2tte_map_size(level);
 
 	for (unsigned int i = 0U; i < S2TTES_PER_S2TT; i++) {
@@ -891,12 +820,15 @@
 void s2tt_init_assigned_ns(unsigned long *s2tt, unsigned long attrs,
 			   unsigned long pa, long level)
 {
+	assert(level >= S2TT_MIN_BLOCK_LEVEL);
+	assert(level <= S2TT_PAGE_LEVEL);
+	assert(s2tt != NULL);
+	assert(s2tte_is_addr_lvl_aligned(pa, level));
+
 	const unsigned long map_size = s2tte_map_size(level);
 
-	assert(addr_is_level_aligned(pa, level));
-
 	for (unsigned int i = 0U; i < S2TTES_PER_S2TT; i++) {
-		unsigned long s2tte = attrs & S2TTE_NS_ATTR_MASK;
+		unsigned long s2tte = attrs & S2TTE_NS_ATTR_HOST_MASK;
 
 		s2tt[i] = s2tte_create_assigned_ns(s2tte | pa, level);
 		pa += map_size;
@@ -915,7 +847,7 @@
  */
 static bool s2tte_has_pa(unsigned long s2tte, long level)
 {
-	unsigned long desc_type = s2tte & DESC_TYPE_MASK;
+	unsigned long desc_type = s2tte & S2TT_DESC_TYPE_MASK;
 
 	return ((desc_type != S2TTE_INVALID) ||	/* block, page or table */
 		s2tte_is_assigned_empty(s2tte, level) ||
@@ -934,38 +866,32 @@
 	return s2tte_has_pa(s2tte, level);
 }
 
-/* Returns physical address of a page entry or block */
+/* Returns physical address of a S2TTE */
 unsigned long s2tte_pa(unsigned long s2tte, long level)
 {
+	assert(level <= S2TT_PAGE_LEVEL);
+	assert(level >= S2TT_MIN_STARTING_LEVEL);
+
 	if (!s2tte_has_pa(s2tte, level)) {
 		assert(false);
 	}
 
 	if (s2tte_is_table(s2tte, level)) {
-		return addr_level_mask(s2tte, RTT_PAGE_LEVEL);
+		return addr_level_mask(s2tte, S2TT_PAGE_LEVEL);
 	}
 
 	return addr_level_mask(s2tte, level);
 }
 
-/* Returns physical address of a table entry */
-unsigned long s2tte_pa_table(unsigned long s2tte, long level)
-{
-	(void)level;
-
-	assert(s2tte_is_table(s2tte, level));
-	return addr_level_mask(s2tte, RTT_PAGE_LEVEL);
-}
-
-bool addr_is_level_aligned(unsigned long addr, long level)
+bool s2tte_is_addr_lvl_aligned(unsigned long addr, long level)
 {
 	return (addr == addr_level_mask(addr, level));
 }
 
 typedef bool (*s2tte_type_checker)(unsigned long s2tte);
 
-static bool __table_is_uniform_block(unsigned long *table,
-					s2tte_type_checker s2tte_is_x)
+static bool table_is_uniform_block(unsigned long *table,
+				   s2tte_type_checker s2tte_is_x)
 {
 	for (unsigned int i = 0U; i < S2TTES_PER_S2TT; i++) {
 		unsigned long s2tte = s2tte_read(&table[i]);
@@ -981,42 +907,52 @@
 /*
  * Returns true if all s2ttes in @table are unassigned_empty.
  */
-bool table_is_unassigned_empty_block(unsigned long *table)
+bool s2tt_is_unassigned_empty_block(unsigned long *table)
 {
-	return __table_is_uniform_block(table, s2tte_is_unassigned_empty);
+	assert(table != NULL);
+
+	return table_is_uniform_block(table, s2tte_is_unassigned_empty);
 }
 
 /*
  * Returns true if all s2ttes in @table are unassigned_ram.
  */
-bool table_is_unassigned_ram_block(unsigned long *table)
+bool s2tt_is_unassigned_ram_block(unsigned long *table)
 {
-	return __table_is_uniform_block(table, s2tte_is_unassigned_ram);
+	assert(table != NULL);
+
+	return table_is_uniform_block(table, s2tte_is_unassigned_ram);
 }
 
 /*
  * Returns true if all s2ttes in @table are unassigned_ns
  */
-bool table_is_unassigned_ns_block(unsigned long *table)
+bool s2tt_is_unassigned_ns_block(unsigned long *table)
 {
-	return __table_is_uniform_block(table, s2tte_is_unassigned_ns);
+	assert(table != NULL);
+
+	return table_is_uniform_block(table, s2tte_is_unassigned_ns);
 }
 
 /*
  * Returns true if all s2ttes in @table are unassigned_destroyed
  */
-bool table_is_unassigned_destroyed_block(unsigned long *table)
+bool s2tt_is_unassigned_destroyed_block(unsigned long *table)
 {
-	return __table_is_uniform_block(table, s2tte_is_unassigned_destroyed);
+	assert(table != NULL);
+
+	return table_is_uniform_block(table, s2tte_is_unassigned_destroyed);
 }
 
 typedef bool (*s2tte_type_level_checker)(unsigned long s2tte, long level);
 
-static bool __table_maps_block(unsigned long *table,
-			       long level,
-			       s2tte_type_level_checker s2tte_is_x,
-			       bool check_ns_attrs)
+static bool table_maps_block(unsigned long *table,
+			     long level,
+			     s2tte_type_level_checker s2tte_is_x,
+			     bool check_ns_attrs)
 {
+	assert(table != NULL);
+
 	unsigned long base_pa;
 	unsigned long map_size = s2tte_map_size(level);
 	unsigned long s2tte = s2tte_read(&table[0]);
@@ -1027,7 +963,7 @@
 	}
 
 	base_pa = s2tte_pa(s2tte, level);
-	if (!addr_is_level_aligned(base_pa, level - 1L)) {
+	if (!s2tte_is_addr_lvl_aligned(base_pa, level - 1L)) {
 		return false;
 	}
 
@@ -1045,13 +981,14 @@
 		}
 
 		if (check_ns_attrs) {
-			unsigned long ns_attrs = s2tte & S2TTE_NS_ATTR_MASK;
+			unsigned long ns_attrs =
+					s2tte & S2TTE_NS_ATTR_HOST_MASK;
 
 			/*
 			 * We match all the attributes in the S2TTE
 			 * except for the AF bit.
 			 */
-			if ((s2tte & S2TTE_NS_ATTR_MASK) != ns_attrs) {
+			if ((s2tte & S2TTE_NS_ATTR_HOST_MASK) != ns_attrs) {
 				return false;
 			}
 		}
@@ -1064,9 +1001,9 @@
  * Returns true if all s2ttes are assigned_empty
  * and refer to a contiguous block of granules aligned to @level - 1.
  */
-bool table_maps_assigned_empty_block(unsigned long *table, long level)
+bool s2tt_maps_assigned_empty_block(unsigned long *table, long level)
 {
-	return __table_maps_block(table, level, s2tte_is_assigned_empty,
+	return table_maps_block(table, level, s2tte_is_assigned_empty,
 				  false);
 }
 
@@ -1074,9 +1011,9 @@
  * Returns true if all s2ttes are assigned_ram and
  * refer to a contiguous block of granules aligned to @level - 1.
  */
-bool table_maps_assigned_ram_block(unsigned long *table, long level)
+bool s2tt_maps_assigned_ram_block(unsigned long *table, long level)
 {
-	return __table_maps_block(table, level, s2tte_is_assigned_ram, false);
+	return table_maps_block(table, level, s2tte_is_assigned_ram, false);
 }
 
 /*
@@ -1087,18 +1024,18 @@
  *
  * @pre: @table maps IPA outside PAR.
  */
-bool table_maps_assigned_ns_block(unsigned long *table, long level)
+bool s2tt_maps_assigned_ns_block(unsigned long *table, long level)
 {
-	return __table_maps_block(table, level, s2tte_is_assigned_ns, true);
+	return table_maps_block(table, level, s2tte_is_assigned_ns, true);
 }
 
 /*
  * Returns true if all s2ttes are assigned_destroyed and
  * refer to a contiguous block of granules aligned to @level - 1.
  */
-bool table_maps_assigned_destroyed_block(unsigned long *table, long level)
+bool s2tt_maps_assigned_destroyed_block(unsigned long *table, long level)
 {
-	return __table_maps_block(table, level, s2tte_is_assigned_destroyed,
+	return table_maps_block(table, level, s2tte_is_assigned_destroyed,
 				  false);
 }
 
@@ -1121,10 +1058,16 @@
  *   first entry in the next table.
  * - Otherwise, the address of the first live entry in @s2tt
  */
-unsigned long skip_non_live_entries(unsigned long addr,
-				    unsigned long *s2tt,
-				    const struct rtt_walk *wi)
+unsigned long s2tt_skip_non_live_entries(unsigned long addr,
+					 unsigned long *table,
+					 const struct s2tt_walk *wi)
 {
+	assert(table != NULL);
+	assert(wi != NULL);
+	assert(wi->index <= S2TTES_PER_S2TT);
+	assert(wi->last_level >= S2TT_MIN_STARTING_LEVEL);
+	assert(wi->last_level <= S2TT_PAGE_LEVEL);
+
 	unsigned long i, index = wi->index;
 	long level = wi->last_level;
 	unsigned long map_size;
@@ -1133,7 +1076,7 @@
 	 * If the entry for the map_addr is live,
 	 * return @addr.
 	 */
-	if (s2tte_is_live(s2tte_read(&s2tt[index]), level)) {
+	if (s2tte_is_live(s2tte_read(&table[index]), level)) {
 		return addr;
 	}
 
@@ -1146,7 +1089,7 @@
 
 	/* Skip the "index" */
 	for (i = index + 1UL; i < S2TTES_PER_S2TT; i++) {
-		unsigned long s2tte = s2tte_read(&s2tt[i]);
+		unsigned long s2tte = s2tte_read(&table[i]);
 
 		if (s2tte_is_live(s2tte, level)) {
 			break;