#error Architecture not supported
 #endif /* arch */
 
+#define PKEY_MASK      (PKEY_DISABLE_ACCESS | PKEY_DISABLE_WRITE)
+
+static inline u64 set_pkey_bits(u64 reg, int pkey, u64 flags)
+{
+       u32 shift = pkey_bit_position(pkey);
+       /* mask out bits from pkey in old value */
+       reg &= ~((u64)PKEY_MASK << shift);
+       /* OR in new bits for pkey */
+       reg |= (flags & PKEY_MASK) << shift;
+       return reg;
+}
+
+static inline u64 get_pkey_bits(u64 reg, int pkey)
+{
+       u32 shift = pkey_bit_position(pkey);
+       /*
+        * shift down the relevant bits to the lowest two, then
+        * mask off all the other higher bits
+        */
+       return ((reg >> shift) & PKEY_MASK);
+}
+
 extern u64 shadow_pkey_reg;
 
 static inline u64 _read_pkey_reg(int line)
 
 
 static u32 hw_pkey_get(int pkey, unsigned long flags)
 {
-       u32 mask = (PKEY_DISABLE_ACCESS|PKEY_DISABLE_WRITE);
        u64 pkey_reg = __read_pkey_reg();
-       u64 shifted_pkey_reg;
-       u32 masked_pkey_reg;
 
        dprintf1("%s(pkey=%d, flags=%lx) = %x / %d\n",
                        __func__, pkey, flags, 0, 0);
        dprintf2("%s() raw pkey_reg: %016llx\n", __func__, pkey_reg);
 
-       shifted_pkey_reg = (pkey_reg >> (pkey * PKEY_BITS_PER_PKEY));
-       dprintf2("%s() shifted_pkey_reg: %016llx\n", __func__,
-                       shifted_pkey_reg);
-       masked_pkey_reg = shifted_pkey_reg & mask;
-       dprintf2("%s() masked  pkey_reg: %x\n", __func__, masked_pkey_reg);
-       /*
-        * shift down the relevant bits to the lowest two, then
-        * mask off all the other high bits.
-        */
-       return masked_pkey_reg;
+       return (u32) get_pkey_bits(pkey_reg, pkey);
 }
 
 static int hw_pkey_set(int pkey, unsigned long rights, unsigned long flags)
        /* make sure that 'rights' only contains the bits we expect: */
        assert(!(rights & ~mask));
 
-       /* copy old pkey_reg */
-       new_pkey_reg = old_pkey_reg;
-       /* mask out bits from pkey in old value: */
-       new_pkey_reg &= ~(mask << (pkey * PKEY_BITS_PER_PKEY));
-       /* OR in new bits for pkey: */
-       new_pkey_reg |= (rights << (pkey * PKEY_BITS_PER_PKEY));
+       /* modify bits accordingly in old pkey_reg and assign it */
+       new_pkey_reg = set_pkey_bits(old_pkey_reg, pkey, rights);
 
        __write_pkey_reg(new_pkey_reg);
 
        ret = hw_pkey_set(pkey, pkey_rights, syscall_flags);
        assert(!ret);
        /* pkey_reg and flags have the same format */
-       shadow_pkey_reg |= flags << (pkey * 2);
+       shadow_pkey_reg = set_pkey_bits(shadow_pkey_reg, pkey, pkey_rights);
        dprintf1("%s(%d) shadow: 0x%016llx\n",
                __func__, pkey, shadow_pkey_reg);
 
        pkey_rights |= flags;
 
        ret = hw_pkey_set(pkey, pkey_rights, 0);
-       shadow_pkey_reg &= ~(flags << (pkey * 2));
+       shadow_pkey_reg = set_pkey_bits(shadow_pkey_reg, pkey, pkey_rights);
        pkey_assert(ret >= 0);
 
        pkey_rights = hw_pkey_get(pkey, syscall_flags);
                        shadow_pkey_reg);
        if (ret) {
                /* clear both the bits: */
-               shadow_pkey_reg &= ~(0x3      << (ret * 2));
+               shadow_pkey_reg = set_pkey_bits(shadow_pkey_reg, ret,
+                                               ~PKEY_MASK);
                dprintf4("%s()::%d, ret: %d pkey_reg: 0x%016llx"
                                " shadow: 0x%016llx\n",
                                __func__,
                 * move the new state in from init_val
                 * (remember, we cheated and init_val == pkey_reg format)
                 */
-               shadow_pkey_reg |=  (init_val << (ret * 2));
+               shadow_pkey_reg = set_pkey_bits(shadow_pkey_reg, ret,
+                                               init_val);
        }
        dprintf4("%s()::%d, ret: %d pkey_reg: 0x%016llx"
                        " shadow: 0x%016llx\n",