#include <palacios/vmm_types.h>
 #include <palacios/vmm_paging.h>
+#include <palacios/vmm_hashtable.h>
 
 
 struct v3_swap_ops {
 
 struct v3_sym_swap_state {
     struct v3_swap_dev devs[256];
+
+    // shadow pointers
+    struct hashtable * shdw_ptr_ht;
 };
 
 
+
 static inline int is_swapped_pte32(pte32_t * pte) {
     return ((pte->present == 0) && (*(uint32_t *)pte != 0));
 }
 
 
 
-
 int v3_init_sym_swap(struct guest_info * info);
 
-addr_t v3_get_swapped_pg_addr(struct guest_info * info, pte32_t * pte);
-
 int v3_register_swap_disk(struct guest_info * info, int dev_index, 
                          struct v3_swap_ops * ops, void * private_data);
 
-int v3_swap_out_notify(struct guest_info * info, int pg_index, int dev_index);
+int v3_swap_in_notify(struct guest_info * info, int pg_index, int dev_index);
+
+
+addr_t v3_get_swapped_pg_addr(struct guest_info * info, pte32_t * shadow_pte, pte32_t * guest_pte);
+
+int v3_swap_flush(struct guest_info * info);
+
 
 #endif
 #endif
 
     uint_t unswapped_pages;
 
 
+    union swap_header * hdr;
+
     uint64_t capacity;
     uint8_t * swap_space;
     addr_t swap_base_addr;
 static inline int get_index_usage(struct swap_state * swap, uint32_t index) {
     int major = index / 8;
     int minor = index % 8;
-    return swap->usage_map[major] & (1 << minor);
+
+    return (swap->usage_map[major] & (1 << minor));
 }
 
 
 
 
 static inline void * get_swap_entry(uint32_t pg_index, void * private_data) {
-    struct swap_state * swap = (struct swap_state *)private_data;
+    struct vm_device * dev = (struct vm_device *)private_data;
+    struct swap_state * swap = (struct swap_state *)(dev->private_data);
     void * pg_addr = NULL;
+    int ret = 0;
 
-    if (get_index_usage(swap, pg_index)) {
+    if ((ret = get_index_usage(swap, pg_index))) {
        // CAREFUL: The index might be offset by 1, because the first 4K is the header
        pg_addr = (void *)(swap->swap_space + (pg_index * 4096));
     }
     uint32_t offset = lba * HD_SECTOR_SIZE;
     uint32_t length = sector_count * HD_SECTOR_SIZE;
 
-    
-    PrintDebug("SymSwap: Reading %d bytes to %p from %p\n", length,
-              buf, (void *)(swap->swap_space + offset));
-    
+  
+    /*  
+       PrintDebug("SymSwap: Reading %d bytes to %p from %p\n", length,
+       buf, (void *)(swap->swap_space + offset));
+    */
+
     if (length % 4096) {
        PrintError("Swapping in length that is not a page multiple\n");
     }
 
     swap->unswapped_pages += (length / 4096);
 
-    set_index_usage(swap, get_swap_index_from_offset(offset), 0);
-
-
-    // Notify the shadow paging layer
-
-    PrintDebug("Swapped in %d pages\n", length / 4096);
+    if ((swap->active == 1) && (offset != 0)) {
+       int i = 0;
+       // Notify the shadow paging layer
+       PrintDebug("Swapped in %d pages\n", length / 4096);
+       for (i = 0; i < length; i += 4096) {
+           set_index_usage(swap, get_swap_index_from_offset(offset + i), 0);
+           v3_swap_in_notify(dev->vm, get_swap_index_from_offset(offset + i), swap->hdr->info.type);
+       }
+    }
 
     return 0;
 }
 
     if ((swap->active == 0) && (offset == 0)) {
        // This is the swap header page
-       union swap_header * hdr = (union swap_header *)buf;
 
        if (length != 4096) {
            PrintError("Initializing Swap space by not writing page multiples. This sucks...\n");
 
        swap->active = 1;
 
-       PrintDebug("Swap Type=%d (magic=%s)\n", hdr->info.type, hdr->magic.magic);
+       PrintDebug("Swap Type=%d (magic=%s)\n", swap->hdr->info.type, swap->hdr->magic.magic);
 
-       if (v3_register_swap_disk(dev->vm, hdr->info.type, &swap_ops, dev) == -1) {
+       if (v3_register_swap_disk(dev->vm, swap->hdr->info.type, &swap_ops, dev) == -1) {
            PrintError("Error registering symbiotic swap disk\n");
            return -1;
        }
-
     }
 
     memcpy(swap->swap_space + offset, buf, length);
 
     swap->swapped_pages += (length / 4096);
 
-    set_index_usage(swap, get_swap_index_from_offset(offset), 1);
+    if ((swap->active == 1) && (offset != 0)) {
+       int i = 0;
+       PrintDebug("Swapped out %d pages\n", length / 4096);
 
-    PrintDebug("Swapped out %d pages\n", length / 4096);
+       for (i = 0; i < length; i += 4096) {
+           set_index_usage(swap, get_swap_index_from_offset(offset + i), 1);
+       }
+    }
 
     return 0;
 }
     swap->unswapped_pages = 0;
 
     swap->active = 0;
+    swap->hdr = (union swap_header *)swap;
+
 
     swap->swap_base_addr = (addr_t)V3_AllocPages(swap->capacity / 4096);
     swap->swap_space = (uint8_t *)V3_VAddr((void *)(swap->swap_base_addr));
     return 0;
 }
 
-
-
 device_register("SYM_SWAP", swap_init)
 
 #include <palacios/vmm_hypercall.h>
 #include <palacios/vmm_dev_mgr.h>
 
+#ifdef CONFIG_SYMBIOTIC_SWAP
+#include <palacios/vmm_sym_swap.h>
+#endif
+
 
 #include <devices/generic.h>
 #include <devices/ide.h>
 #include <devices/video.h>
 
 
+
+
 #include <palacios/vmm_host_events.h>
 
-#define USE_GENERIC 1
 
 
 #include <palacios/vmm_socket.h>
     
     v3_init_hypercall_map(info);
 
+#ifdef CONFIG_SYMBIOTIC_SWAP
+    v3_init_sym_swap(info);
+#endif
+
+
 
     // Initialize the memory map
     v3_init_shadow_map(info);
     v3_create_device(info, "IOAPIC", "LAPIC");
     v3_create_device(info, "VMNET", NULL);
     
-    int use_generic = USE_GENERIC;
+
 
     if (config_ptr->enable_pci == 1) {
        struct ide_cfg ide_config = {"PCI", "PIIX3"};
 
 
 
-    if (use_generic) {
+#ifdef CONFIG_GENERIC
        configure_generic(info, config_ptr);
-    }
+#endif
 
     // This should go last because it requires information about the Harddrives
     v3_create_device(info, "NVRAM", "IDE");
 
 
 
-
+#ifdef CONFIG_GENERIC
 static int configure_generic(struct guest_info * info, struct v3_vm_config * config_ptr) {
     PrintDebug("Creating Generic Device\n");
     v3_create_device(info, "GENERIC", NULL);
     
     return 0;
 }
+#endif
 
     shadow_cr3->pwt = guest_cr3->pwt;
     shadow_cr3->pcd = guest_cr3->pcd;
   
+#ifdef CONFIG_SYMBIOTIC_SWAP
+    v3_swap_flush(info);
+#endif
+
     return 0;
 }
 
        if (is_swapped_pte32(guest_pte)) {
            PrintError("Page fault on swapped out page (pte=%x)\n", *(uint32_t *)guest_pte);
 
-           if (inject_guest_pf(info, fault_addr, error_code) == -1) {
-               PrintError("Could not inject guest page fault\n");
-               return -1;
+           addr_t swp_pg_addr = v3_get_swapped_pg_addr(info, shadow_pte, guest_pte);
+
+           if (swp_pg_addr == 0) {
+               if (inject_guest_pf(info, fault_addr, error_code) == -1) {
+                   PrintError("Could not inject guest page fault\n");
+                   return -1;
+               }
+           } else {
+               /* 
+                *  Setup shadow paging state
+                */
+               
+               /* We need some way to check permissions.... */
+
+               shadow_pte->accessed = 1;
+               shadow_pte->writable = 1;
+               shadow_pte->write_through = 0;
+               shadow_pte->cache_disable = 0;
+               shadow_pte->global_page = 0;
+               shadow_pte->user_page = 1;
+               shadow_pte->present = 1;
+               
+               shadow_pte->page_base_addr = swp_pg_addr;
            }
        } else {
            if (inject_guest_pf(info, fault_addr, error_code) == -1) {
 
 
 
 #include <palacios/vmm_sym_swap.h>
+#include <palacios/vmm_list.h>
+
+
+// This is a hack and 32 bit linux specific.... need to fix...
+struct swap_pte {
+    uint32_t present    : 1;
+    uint32_t dev_index  : 8;
+    uint32_t pg_index   : 23;
+};
+
+
+struct shadow_pointer {
+    uint32_t pg_index;
+    uint32_t dev_index;
+
+    pte32_t * shadow_pte;
+    
+    addr_t guest_pte;
+    
+    struct list_head node;
+};
+
+
+static uint_t swap_hash_fn(addr_t key) {
+    return v3_hash_long(key, 32);
+}
+
+
+static int swap_eq_fn(addr_t key1, addr_t key2) {
+    return (key1 == key2);
+}
 
-// this is hardcoded in linux, but we should expose it via a sym interface
-#define SWAP_DEV_SHIFT 5 
 
-int v3_init_sym_swap(struct guest_info * info) {
 
+static inline uint32_t get_pg_index(pte32_t * pte) {
+    return ((struct swap_pte *)pte)->pg_index;
+}
+
+
+static inline uint32_t get_dev_index(pte32_t * pte) {
+    return ((struct swap_pte *)pte)->dev_index;
+}
+
+
+
+int v3_init_sym_swap(struct guest_info * info) {
     memset(&(info->swap_state), 0, sizeof(struct v3_sym_swap_state));
+    info->swap_state.shdw_ptr_ht = v3_create_htable(0, swap_hash_fn, swap_eq_fn);
+
+    PrintDebug("Initialized Symbiotic Swap\n");
 
     return 0;
 }
 }
 
 
-addr_t v3_get_swapped_pg_addr(struct guest_info * info, pte32_t * pte) {
+
+
+int v3_swap_in_notify(struct guest_info * info, int pg_index, int dev_index) {
+    struct list_head * shdw_ptr_list = NULL;
     struct v3_sym_swap_state * swap_state = &(info->swap_state);
-    uint32_t dev_index = *(uint32_t *)pte & ((1 << SWAP_DEV_SHIFT) - 1);
-    uint32_t pg_index = (*(uint32_t *)pte) >> SWAP_DEV_SHIFT;
-    struct v3_swap_dev * swp_dev = &(swap_state->devs[dev_index]);
+    struct shadow_pointer * tmp_shdw_ptr = NULL;
+    struct shadow_pointer * shdw_ptr = NULL;
+    struct swap_pte guest_pte = {0, dev_index, pg_index};
 
-    return (addr_t)swp_dev->ops->get_swap_entry(pg_index, swp_dev->private_data);
+    shdw_ptr_list = (struct list_head * )v3_htable_search(swap_state->shdw_ptr_ht, *(addr_t *)&(guest_pte));
+
+    if (shdw_ptr_list == NULL) {
+       return 0;
+    }
+
+    list_for_each_entry_safe(shdw_ptr, tmp_shdw_ptr, shdw_ptr_list, node) {
+       if ((shdw_ptr->pg_index == pg_index) &&
+           (shdw_ptr->dev_index == dev_index)) {
+
+           // Trigger faults for next shadow access
+           shdw_ptr->shadow_pte->present = 0;
+
+           // Delete entry from list
+           list_del(&(shdw_ptr->node));
+           V3_Free(shdw_ptr);
+       }
+    }
+
+    return 0;
 }
 
 
 
-int v3_swap_out_notify(struct guest_info * info, int pg_index, int dev_index) {
-    return -1;
+int v3_swap_flush(struct guest_info * info) {
+    struct v3_sym_swap_state * swap_state = &(info->swap_state);
+    struct hashtable_iter * ht_iter = v3_create_htable_iter(swap_state->shdw_ptr_ht);
+
+    PrintDebug("Flushing Symbiotic Swap table\n");
+
+    while (ht_iter->entry) {
+       struct shadow_pointer * tmp_shdw_ptr = NULL;
+       struct shadow_pointer * shdw_ptr = NULL;
+       struct list_head * shdw_ptr_list = (struct list_head *)v3_htable_get_iter_value(ht_iter);
+
+       // delete all swapped entries
+       // we can leave the list_head structures and reuse them for the next round
+       
+       list_for_each_entry_safe(shdw_ptr, tmp_shdw_ptr, shdw_ptr_list, node) {
+           // Trigger faults for next shadow access
+           shdw_ptr->shadow_pte->present = 0;
+           
+           // Delete entry from list
+           list_del(&(shdw_ptr->node));
+           V3_Free(shdw_ptr);      
+       }
+
+       v3_htable_iter_advance(ht_iter);
+    }
+
+    return 0;
+}
+
+
+addr_t v3_get_swapped_pg_addr(struct guest_info * info, pte32_t * shadow_pte, pte32_t * guest_pte) {
+    struct list_head * shdw_ptr_list = NULL;
+    struct v3_sym_swap_state * swap_state = &(info->swap_state);
+    struct shadow_pointer * shdw_ptr = NULL;
+    void * swp_page_ptr = NULL;
+    int dev_index = get_dev_index(guest_pte);
+    struct v3_swap_dev * swp_dev = &(swap_state->devs[dev_index]);
+
+    if (! swp_dev->present ) {
+       return 0;
+    }
+
+    swp_page_ptr = swp_dev->ops->get_swap_entry(get_pg_index(guest_pte), swp_dev->private_data);
+
+    if (swp_page_ptr == NULL) {
+       PrintError("Swapped out page not found on swap device\n");
+       return 0;
+    }
+
+    shdw_ptr_list = (struct list_head *)v3_htable_search(swap_state->shdw_ptr_ht, (addr_t)*(uint32_t *)guest_pte);
+
+    if (shdw_ptr_list == NULL) {
+       shdw_ptr_list = (struct list_head *)V3_Malloc(sizeof(struct list_head *));
+       INIT_LIST_HEAD(shdw_ptr_list);
+    }
+
+    shdw_ptr = (struct shadow_pointer *)V3_Malloc(sizeof(struct shadow_pointer));
+
+    shdw_ptr->shadow_pte = shadow_pte;
+    shdw_ptr->guest_pte = *(uint32_t *)guest_pte;
+    shdw_ptr->pg_index = get_pg_index(guest_pte);
+    shdw_ptr->dev_index = get_dev_index(guest_pte);
+
+    // We don't check for conflicts, because it should not happen...
+    list_add(&(shdw_ptr->node), shdw_ptr_list);
+
+    return PAGE_BASE_ADDR((addr_t)V3_PAddr(swp_page_ptr));
 }