1 files changed, 21 insertions, 14 deletions

diff --git a/lib/sbi/sbi_domain_context.c b/lib/sbi/sbi_domain_context.c
index 407c0d5c..2b19fcc7 100644
--- a/lib/sbi/sbi_domain_context.c
+++ b/lib/sbi/sbi_domain_context.c
@@ -53,31 +53,30 @@ struct hart_context {
 	bool initialized;
 };
 
-struct domain_context_priv {
-	/** Contexts for possible HARTs indexed by hartindex */
-	struct hart_context *hartindex_to_context_table[SBI_HARTMASK_MAX_BITS];
-};
-
-static struct sbi_domain_data dcpriv = {
-	.data_size = sizeof(struct domain_context_priv),
-};
+static struct sbi_domain_data dcpriv;
 
 static inline struct hart_context *hart_context_get(struct sbi_domain *dom,
 						    u32 hartindex)
 {
-	struct domain_context_priv *dcp = sbi_domain_data_ptr(dom, &dcpriv);
+	struct hart_context **dom_hartindex_to_context_table;
+
+	dom_hartindex_to_context_table = sbi_domain_data_ptr(dom, &dcpriv);
+	if (!dom_hartindex_to_context_table || !sbi_hartindex_valid(hartindex))
+		return NULL;
 
-	return (dcp && hartindex < SBI_HARTMASK_MAX_BITS) ?
-		dcp->hartindex_to_context_table[hartindex] : NULL;
+	return dom_hartindex_to_context_table[hartindex];
 }
 
 static void hart_context_set(struct sbi_domain *dom, u32 hartindex,
 			     struct hart_context *hc)
 {
-	struct domain_context_priv *dcp = sbi_domain_data_ptr(dom, &dcpriv);
+	struct hart_context **dom_hartindex_to_context_table;
 
-	if (dcp && hartindex < SBI_HARTMASK_MAX_BITS)
-		dcp->hartindex_to_context_table[hartindex] = hc;
+	dom_hartindex_to_context_table = sbi_domain_data_ptr(dom, &dcpriv);
+	if (!dom_hartindex_to_context_table || !sbi_hartindex_valid(hartindex))
+		return;
+
+	dom_hartindex_to_context_table[hartindex] = hc;
 }
 
 /** Macro to obtain the current hart's context pointer */
@@ -232,6 +231,14 @@ int sbi_domain_context_exit(void)
 
 int sbi_domain_context_init(void)
 {
+	/**
+	 * Allocate per-domain and per-hart context data.
+	 * The data type is "struct hart_context **" whose memory space will be
+	 * dynamically allocated by domain_setup_data_one(). Calculate needed
+	 * size of memory space here.
+	 */
+	dcpriv.data_size = sizeof(struct hart_context *) * sbi_hart_count();
+
 	return sbi_domain_register_data(&dcpriv);
 }