From 5b8cde66f120717e063ac6a45e0690332c8b5219 Mon Sep 17 00:00:00 2001
From: Jacek Caban <jacek@codeweavers.com>
Date: Mon, 12 Sep 2011 12:29:14 +0200
Subject: [PATCH] vbscript: Make vbscode_t own the memory it uses.

---
 dlls/vbscript/compile.c       | 33 ++++++++++++--
 dlls/vbscript/vbscript.h      | 13 ++++++
 dlls/vbscript/vbscript_main.c | 83 +++++++++++++++++++++++++++++++++++
 3 files changed, 125 insertions(+), 4 deletions(-)

diff --git a/dlls/vbscript/compile.c b/dlls/vbscript/compile.c
index 181f6c1d262..450fa183d2b 100644
--- a/dlls/vbscript/compile.c
+++ b/dlls/vbscript/compile.c
@@ -36,6 +36,23 @@ typedef struct {
 
 static HRESULT compile_expression(compile_ctx_t*,expression_t*);
 
+static inline void *compiler_alloc(vbscode_t *vbscode, size_t size)
+{
+    return vbsheap_alloc(&vbscode->heap, size);
+}
+
+static WCHAR *compiler_alloc_string(vbscode_t *vbscode, const WCHAR *str)
+{
+    size_t size;
+    WCHAR *ret;
+
+    size = (strlenW(str)+1)*sizeof(WCHAR);
+    ret = compiler_alloc(vbscode, size);
+    if(ret)
+        memcpy(ret, str, size);
+    return ret;
+}
+
 static inline instr_t *instr_ptr(compile_ctx_t *ctx, unsigned id)
 {
     assert(id < ctx->instr_cnt);
@@ -75,13 +92,18 @@ static HRESULT push_instr_int(compile_ctx_t *ctx, vbsop_t op, LONG arg)
 
 static HRESULT push_instr_str(compile_ctx_t *ctx, vbsop_t op, const WCHAR *arg)
 {
-    unsigned ret;
+    unsigned instr;
+    WCHAR *str;
 
-    ret = push_instr(ctx, op);
-    if(ret == -1)
+    str = compiler_alloc_string(ctx->code, arg);
+    if(!str)
         return E_OUTOFMEMORY;
 
-    instr_ptr(ctx, ret)->arg1.str = arg;
+    instr = push_instr(ctx, op);
+    if(instr == -1)
+        return E_OUTOFMEMORY;
+
+    instr_ptr(ctx, instr)->arg1.str = str;
     return S_OK;
 }
 
@@ -263,6 +285,8 @@ void release_vbscode(vbscode_t *code)
     for(i=0; i < code->bstr_cnt; i++)
         SysFreeString(code->bstr_pool[i]);
 
+    vbsheap_free(&code->heap);
+
     heap_free(code->bstr_pool);
     heap_free(code->source);
     heap_free(code->instrs);
@@ -291,6 +315,7 @@ static vbscode_t *alloc_vbscode(compile_ctx_t *ctx, const WCHAR *source)
 
     ctx->instr_cnt = 0;
     ctx->instr_size = 32;
+    vbsheap_init(&ret->heap);
 
     ret->option_explicit = ctx->parser.option_explicit;
 
diff --git a/dlls/vbscript/vbscript.h b/dlls/vbscript/vbscript.h
index 199ca70e90c..9c047d16c13 100644
--- a/dlls/vbscript/vbscript.h
+++ b/dlls/vbscript/vbscript.h
@@ -31,6 +31,18 @@
 #include "wine/list.h"
 #include "wine/unicode.h"
 
+typedef struct {
+    void **blocks;
+    DWORD block_cnt;
+    DWORD last_block;
+    DWORD offset;
+    struct list custom_blocks;
+} vbsheap_t;
+
+void vbsheap_init(vbsheap_t*) DECLSPEC_HIDDEN;
+void *vbsheap_alloc(vbsheap_t*,size_t) __WINE_ALLOC_SIZE(2) DECLSPEC_HIDDEN;
+void vbsheap_free(vbsheap_t*) DECLSPEC_HIDDEN;
+
 typedef struct _function_t function_t;
 typedef struct _vbscode_t vbscode_t;
 typedef struct _script_ctx_t script_ctx_t;
@@ -122,6 +134,7 @@ struct _vbscode_t {
     BSTR *bstr_pool;
     unsigned bstr_pool_size;
     unsigned bstr_cnt;
+    vbsheap_t heap;
 
     struct list entry;
 };
diff --git a/dlls/vbscript/vbscript_main.c b/dlls/vbscript/vbscript_main.c
index 9d74026792d..15272810773 100644
--- a/dlls/vbscript/vbscript_main.c
+++ b/dlls/vbscript/vbscript_main.c
@@ -31,6 +31,89 @@ DEFINE_GUID(GUID_NULL,0,0,0,0,0,0,0,0,0,0,0);
 
 static HINSTANCE vbscript_hinstance;
 
+#define MIN_BLOCK_SIZE  128
+
+static inline DWORD block_size(DWORD block)
+{
+    return MIN_BLOCK_SIZE << block;
+}
+
+void vbsheap_init(vbsheap_t *heap)
+{
+    memset(heap, 0, sizeof(*heap));
+    list_init(&heap->custom_blocks);
+}
+
+void *vbsheap_alloc(vbsheap_t *heap, size_t size)
+{
+    struct list *list;
+    void *tmp;
+
+    size = (size+3)&~3;
+
+    if(!heap->block_cnt) {
+        if(!heap->blocks) {
+            heap->blocks = heap_alloc(sizeof(void*));
+            if(!heap->blocks)
+                return NULL;
+        }
+
+        tmp = heap_alloc(block_size(0));
+        if(!tmp)
+            return NULL;
+
+        heap->blocks[0] = tmp;
+        heap->block_cnt = 1;
+    }
+
+    if(heap->offset + size <= block_size(heap->last_block)) {
+        tmp = ((BYTE*)heap->blocks[heap->last_block])+heap->offset;
+        heap->offset += size;
+        return tmp;
+    }
+
+    if(size <= block_size(heap->last_block+1)) {
+        if(heap->last_block+1 == heap->block_cnt) {
+            tmp = heap_realloc(heap->blocks, (heap->block_cnt+1)*sizeof(void*));
+            if(!tmp)
+                return NULL;
+
+            heap->blocks = tmp;
+            heap->blocks[heap->block_cnt] = heap_alloc(block_size(heap->block_cnt));
+            if(!heap->blocks[heap->block_cnt])
+                return NULL;
+
+            heap->block_cnt++;
+        }
+
+        heap->last_block++;
+        heap->offset = size;
+        return heap->blocks[heap->last_block];
+    }
+
+    list = heap_alloc(size + sizeof(struct list));
+    if(!list)
+        return NULL;
+
+    list_add_head(&heap->custom_blocks, list);
+    return list+1;
+}
+
+void vbsheap_free(vbsheap_t *heap)
+{
+    struct list *iter;
+    DWORD i;
+
+    while((iter = list_next(&heap->custom_blocks, &heap->custom_blocks))) {
+        list_remove(iter);
+        heap_free(iter);
+    }
+
+    for(i=0; i < heap->block_cnt; i++)
+        heap_free(heap->blocks[i]);
+    heap_free(heap->blocks);
+}
+
 static HRESULT WINAPI ClassFactory_QueryInterface(IClassFactory *iface, REFIID riid, void **ppv)
 {
     *ppv = NULL;