diff --git a/include/uc_priv.h b/include/uc_priv.h
index 14ed07bd..c2b7baa0 100644
--- a/include/uc_priv.h
+++ b/include/uc_priv.h
@@ -389,6 +389,7 @@ struct uc_struct {
     uint64_t qemu_real_host_page_size;
     int qemu_icache_linesize;
     /* ARCH_REGS_STORAGE_SIZE */
+    uc_context_content context_content;
     int cpu_context_size;
     uint64_t next_pc; // save next PC for some special cases
     bool hook_insert; // insert new hook at begin of the hook list (append by
@@ -419,6 +420,7 @@ struct uc_context {
     size_t context_size; // size of the real internal context structure
     uc_mode mode;        // the mode of this context
     uc_arch arch;        // the arch of this context
+    int snapshot_level;  // the memory snapshot level to restore
     char data[0];        // context
 };
 
diff --git a/include/unicorn/unicorn.h b/include/unicorn/unicorn.h
index f2030546..85a1ade8 100644
--- a/include/unicorn/unicorn.h
+++ b/include/unicorn/unicorn.h
@@ -581,6 +581,9 @@ typedef enum uc_control_type {
     // Write: @args = (uint32_t)
     // Read: @args = (uint32_t*)
     UC_CTL_TCG_BUFFER_SIZE,
+    // controle if context_save/restore should work with snapshots
+    // Write: @args = (int)
+    UC_CTL_CONTEXT_MODE,
 } uc_control_type;
 
 /*
@@ -662,6 +665,7 @@ See sample_ctl.c for a detailed example.
     uc_ctl(uc, UC_CTL_READ(UC_CTL_TCG_BUFFER_SIZE, 1), (size))
 #define uc_ctl_set_tcg_buffer_size(uc, size)                                   \
     uc_ctl(uc, UC_CTL_WRITE(UC_CTL_TCG_BUFFER_SIZE, 1), (size))
+#define uc_ctl_context_mode(uc, mode) uc_ctl(uc, UC_CTL_WRITE(UC_CTL_CONTEXT_MODE, 1), (mode))
 
 // Opaque storage for CPU context, used with uc_context_*()
 struct uc_context;
@@ -1015,6 +1019,11 @@ struct uc_tlb_entry {
     uc_prot perms;
 };
 
+typedef enum uc_context_content {
+    UC_CTL_CONTEXT_CPU = 1,
+    UC_CTL_CONTEXT_MEMORY = 2,
+} uc_context_content;
+
 /*
  Map memory in for emulation.
  This API adds a memory region that can be used by emulation.
@@ -1345,12 +1354,6 @@ size_t uc_context_size(uc_engine *uc);
 UNICORN_EXPORT
 uc_err uc_context_free(uc_context *context);
 
-UNICORN_EXPORT
-uc_err uc_snapshot(uc_engine *uc);
-
-UNICORN_EXPORT
-uc_err uc_restore_latest_snapshot(uc_engine *uc);
-
 #ifdef __cplusplus
 }
 #endif
diff --git a/tests/unit/test_mem.c b/tests/unit/test_mem.c
index d25b3448..8f6bcf4f 100644
--- a/tests/unit/test_mem.c
+++ b/tests/unit/test_mem.c
@@ -278,30 +278,34 @@ static void test_mem_protect_mmio(void)
 static void test_snapshot(void)
 {
     uc_engine *uc;
+    uc_context *c0, *c1;
     uint32_t mem;
     // mov eax, [0x2020]; inc eax; mov [0x2020], eax
     char code[] = "\xa1\x20\x20\x00\x00\x00\x00\x00\x00\xff\xc0\xa3\x20\x20\x00"
                   "\x00\x00\x00\x00\x00";
 
     OK(uc_open(UC_ARCH_X86, UC_MODE_64, &uc));
+    OK(uc_context_alloc(uc, &c0));
+    OK(uc_context_alloc(uc, &c1));
+    OK(uc_ctl_context_mode(uc, UC_CTL_CONTEXT_MEMORY));
     OK(uc_mem_map(uc, 0x1000, 0x1000, UC_PROT_ALL));
     OK(uc_mem_write(uc, 0x1000, code, sizeof(code) - 1));
 
     OK(uc_mem_map(uc, 0x2000, 0x1000, UC_PROT_ALL));
-    OK(uc_snapshot(uc));
+    OK(uc_context_save(uc, c0));
 
     OK(uc_emu_start(uc, 0x1000, 0x1000 + sizeof(code) - 1, 0, 0));
     OK(uc_mem_read(uc, 0x2020, &mem, sizeof(mem)));
     TEST_CHECK(mem == 1);
-    OK(uc_snapshot(uc));
+    OK(uc_context_save(uc, c1));
     OK(uc_emu_start(uc, 0x1000, 0x1000 + sizeof(code) - 1, 0, 0));
     OK(uc_mem_read(uc, 0x2020, &mem, sizeof(mem)));
     TEST_CHECK(mem == 2);
-    OK(uc_restore_latest_snapshot(uc));
+    OK(uc_context_restore(uc, c1));
     //TODO check mem
     OK(uc_mem_read(uc, 0x2020, &mem, sizeof(mem)));
     TEST_CHECK(mem == 1);
-    OK(uc_restore_latest_snapshot(uc));
+    OK(uc_context_restore(uc, c0));
     OK(uc_mem_read(uc, 0x2020, &mem, sizeof(mem)));
     TEST_CHECK(mem == 0);
     //TODO check mem
@@ -318,7 +322,7 @@ static void test_context_snapshot(void)
     uint64_t tmp = 1;
 
     OK(uc_open(UC_ARCH_X86, UC_MODE_64, &uc));
-    OK(uc_ctl_context_use_snapshots(uc, 1));
+    OK(uc_ctl_context_mode(uc, UC_CTL_CONTEXT_MEMORY|UC_CTL_CONTEXT_CPU));
     OK(uc_mem_map(uc, 0x1000, 0x1000, UC_PROT_ALL));
     OK(uc_context_alloc(uc, &ctx));
     OK(uc_context_save(uc, ctx));
diff --git a/uc.c b/uc.c
index 184aba32..90bc779e 100644
--- a/uc.c
+++ b/uc.c
@@ -31,6 +31,8 @@
 #include "qemu-common.h"
 
 static void clear_deleted_hooks(uc_engine *uc);
+static uc_err uc_snapshot(uc_engine *uc);
+static uc_err uc_restore_latest_snapshot(uc_engine *uc);
 
 static void *hook_insert(struct list *l, struct hook *h)
 {
@@ -256,6 +258,8 @@ static uc_err uc_init_engine(uc_engine *uc)
         uc->reg_reset(uc);
     }
 
+    uc->context_content = UC_CTL_CONTEXT_CPU;
+
     uc->init_done = true;
 
     return UC_ERR_OK;
@@ -1985,13 +1989,26 @@ UNICORN_EXPORT
 uc_err uc_context_save(uc_engine *uc, uc_context *context)
 {
     UC_INIT(uc);
+    uc_err ret = UC_ERR_OK;
 
-    if (!uc->context_save) {
-        memcpy(context->data, uc->cpu->env_ptr, context->context_size);
-        return UC_ERR_OK;
-    } else {
-        return uc->context_save(uc, context);
+    if (uc->context_content & UC_CTL_CONTEXT_MEMORY) {
+        ret = uc_snapshot(uc);
+        if (ret != UC_ERR_OK) {
+            return ret;
+        }
     }
+
+    context->snapshot_level = uc->snapshot_level;
+
+    if (uc->context_content & UC_CTL_CONTEXT_CPU) {
+        if (!uc->context_save) {
+            memcpy(context->data, uc->cpu->env_ptr, context->context_size);
+            return UC_ERR_OK;
+        } else {
+            return uc->context_save(uc, context);
+        }
+    }
+    return ret;
 }
 
 // Keep in mind that we don't a uc_engine when r/w the registers of a context.
@@ -2237,13 +2254,26 @@ UNICORN_EXPORT
 uc_err uc_context_restore(uc_engine *uc, uc_context *context)
 {
     UC_INIT(uc);
+    uc_err ret;
 
-    if (!uc->context_restore) {
-        memcpy(uc->cpu->env_ptr, context->data, context->context_size);
-        return UC_ERR_OK;
-    } else {
-        return uc->context_restore(uc, context);
+    if (uc->context_content & UC_CTL_CONTEXT_MEMORY) {
+        uc->snapshot_level = context->snapshot_level;
+        ret = uc_restore_latest_snapshot(uc);
+        if (ret != UC_ERR_OK) {
+            return ret;
+        }
+        uc_snapshot(uc);
     }
+
+    if (uc->context_content & UC_CTL_CONTEXT_CPU) {
+        if (!uc->context_restore) {
+            memcpy(uc->cpu->env_ptr, context->data, context->context_size);
+            return UC_ERR_OK;
+        } else {
+            return uc->context_restore(uc, context);
+        }
+    }
+    return UC_ERR_OK;
 }
 
 UNICORN_EXPORT
@@ -2588,6 +2618,19 @@ uc_err uc_ctl(uc_engine *uc, uc_control_type control, ...)
         break;
     }
 
+    case UC_CTL_CONTEXT_MODE:
+
+        UC_INIT(uc);
+
+        if (rw == UC_CTL_IO_WRITE) {
+            int mode = va_arg(args, int);
+            uc->context_content = mode;
+            err = UC_ERR_OK;
+        } else {
+            err = UC_ERR_ARG;
+        }
+        break;
+
     default:
         err = UC_ERR_ARG;
         break;
@@ -2598,8 +2641,7 @@ uc_err uc_ctl(uc_engine *uc, uc_control_type control, ...)
     return err;
 }
 
-UNICORN_EXPORT
-uc_err uc_snapshot(struct uc_struct *uc)
+static uc_err uc_snapshot(struct uc_struct *uc)
 {
     if (uc->snapshot_level == INT32_MAX) {
         return UC_ERR_RESOURCE;
@@ -2608,8 +2650,7 @@ uc_err uc_snapshot(struct uc_struct *uc)
     return UC_ERR_OK;
 }
 
-UNICORN_EXPORT
-uc_err uc_restore_latest_snapshot(struct uc_struct *uc)
+static uc_err uc_restore_latest_snapshot(struct uc_struct *uc)
 {
     MemoryRegion *subregion, *subregion_next;