3 files changed, 29 insertions, 4 deletions

diff --git a/src/gallium/drivers/trace/tr_context.c b/src/gallium/drivers/trace/tr_context.c
index a416628562..2ad5ca4998 100644
--- a/src/gallium/drivers/trace/tr_context.c
+++ b/src/gallium/drivers/trace/tr_context.c
@@ -582,7 +582,7 @@ trace_context_create_fs_state(struct pipe_context *_pipe,
 
    trace_dump_call_end();
 
-   result = trace_shader_create(tr_ctx, state, result);
+   result = trace_shader_create(tr_ctx, state, result, TRACE_SHADER_FRAGMENT);
 
    return result;
 }
@@ -654,7 +654,7 @@ trace_context_create_vs_state(struct pipe_context *_pipe,
 
    trace_dump_call_end();
 
-   result = trace_shader_create(tr_ctx, state, result);
+   result = trace_shader_create(tr_ctx, state, result, TRACE_SHADER_VERTEX);
 
    return result;
 }
diff --git a/src/gallium/drivers/trace/tr_state.c b/src/gallium/drivers/trace/tr_state.c
index c9ce63a0ed..d8c11640bf 100644
--- a/src/gallium/drivers/trace/tr_state.c
+++ b/src/gallium/drivers/trace/tr_state.c
@@ -27,13 +27,18 @@
 #include "util/u_memory.h"
 #include "util/u_simple_list.h"
 
+#include "tgsi/tgsi_parse.h"
+
 struct trace_shader * trace_shader_create(struct trace_context *tr_ctx,
                                           const struct pipe_shader_state *state,
-                                          void *result)
+                                          void *result,
+                                          enum trace_shader_type type)
 {
    struct trace_shader *tr_shdr = CALLOC_STRUCT(trace_shader);
 
    tr_shdr->state = result;
+   tr_shdr->type = type;
+   tr_shdr->tokens = tgsi_dup_tokens(state->tokens);
 
    /* works on context as well */
    trace_screen_add_to_list(tr_ctx, shaders, tr_shdr);
@@ -46,5 +51,16 @@ void trace_shader_destroy(struct trace_context *tr_ctx,
 {
    trace_screen_remove_from_list(tr_ctx, shaders, tr_shdr);
 
+   if (tr_shdr->replaced) {
+      if (tr_shdr->type == TRACE_SHADER_FRAGMENT)
+         tr_ctx->pipe->delete_fs_state(tr_ctx->pipe, tr_shdr->replaced);
+      else if (tr_shdr->type == TRACE_SHADER_VERTEX)
+         tr_ctx->pipe->delete_vs_state(tr_ctx->pipe, tr_shdr->replaced);
+      else
+         assert(0);
+   }
+
+   FREE(tr_shdr->replaced_tokens);
+   FREE(tr_shdr->tokens);
    FREE(tr_shdr);
 }
diff --git a/src/gallium/drivers/trace/tr_state.h b/src/gallium/drivers/trace/tr_state.h
index bca615c50e..1c16042ee5 100644
--- a/src/gallium/drivers/trace/tr_state.h
+++ b/src/gallium/drivers/trace/tr_state.h
@@ -29,10 +29,18 @@
 
 struct tgsi_token;
 
+enum trace_shader_type {
+   TRACE_SHADER_FRAGMENT = 0,
+   TRACE_SHADER_VERTEX   = 1,
+   TRACE_SHADER_GEOMETRY = 2,
+};
+
 struct trace_shader
 {
    struct tr_list list;
 
+   enum trace_shader_type type;
+
    void *state;
    void *replaced;
 
@@ -51,7 +59,8 @@ trace_shader(void *state)
 
 struct trace_shader * trace_shader_create(struct trace_context *tr_ctx,
                                           const struct pipe_shader_state *state,
-                                          void *result);
+                                          void *result,
+                                          enum trace_shader_type type);
 
 void trace_shader_destroy(struct trace_context *tr_ctx,
                           struct trace_shader *tr_shdr);