diff --git a/src/opengx.h b/src/opengx.h
index bbfa0f4..a687bb7 100644
--- a/src/opengx.h
+++ b/src/opengx.h
@@ -153,6 +153,10 @@ typedef void (*OgxSetupDrawCb)(GLuint shader, const OgxDrawData *draw_data,
                                void *user_data);
 void ogx_shader_set_setup_draw_cb(GLuint shader,
                                   OgxSetupDrawCb setup_draw, void *data);
+
+/* These can be called from the setup_draw callback */
+void ogx_shader_setup_attribute_array(int index, uint8_t gx_attr,
+                                      const OgxDrawData *draw_data);
 void *ogx_shader_get_data(GLuint shader);
 
 #ifdef __cplusplus
diff --git a/src/shader.c b/src/shader.c
index 0e8b7bb..d39e385 100644
--- a/src/shader.c
+++ b/src/shader.c
@@ -518,12 +518,12 @@ void glVertexAttribPointer(GLuint index, GLint size, GLenum type,
     OgxVertexAttribState *v = &s_state.vertex_attribs[index];
     if (!v) return;
 
-    v->size = size;
-    v->type = type;
-    v->normalized = normalized;
-    assert(stride < (1 << (8 * sizeof(v->stride))));
-    v->stride = stride;
-    v->pointer = pointer;
+    v->array.size = size;
+    v->array.type = type;
+    v->array.normalized = normalized;
+    assert(stride < (1 << (8 * sizeof(v->array.stride))));
+    v->array.stride = stride;
+    v->array.pointer = pointer;
 }
 
 void ogx_shader_register_program_processor(const OgxProgramProcessor *processor)
diff --git a/src/shader.h b/src/shader.h
index 80b3c8a..9cf14f8 100644
--- a/src/shader.h
+++ b/src/shader.h
@@ -43,23 +43,21 @@ typedef struct {
     OgxProgram *programs;
 
     struct _OgxVertexAttribState {
-        /* Array-related fields (used when array_enabled is true) */
         unsigned array_enabled : 1;
-        unsigned normalized : 1;
-        uint8_t stride;
-        GLenum type;
-        GLint size;
-        const void *pointer;
-
-        /* Data fields (used when array_enabled is false) */
-        union {
-            GLfloat vec4f[4];
-            GLuint vec4ui[4];
-            GLint vec4i[4];
-        } d;
+        /* Used when array_enabled is true: */
+        OgxVertexAttribArray array;
 
         OgxVariable *bound_attribute;
     } vertex_attribs[MAX_VERTEX_ATTRIBS];
+
+    /* Data fields (used when array_enabled is false). We keep these in a
+     * separate struct (well, union) so that we can use consecutive elements as
+     * matrix columns. */
+    union {
+        GLfloat vec4f[4];
+        GLuint vec4ui[4];
+        GLint vec4i[4];
+    } vertex_attrib_data[MAX_VERTEX_ATTRIBS];
 } OgxShaderState;
 
 typedef struct _OgxVertexAttribState OgxVertexAttribState;
diff --git a/src/types.h b/src/types.h
index b910f7a..b25ae75 100644
--- a/src/types.h
+++ b/src/types.h
@@ -52,7 +52,8 @@ typedef float Norm3f[3];
 typedef float Tex2f[2];
 
 typedef struct _OgxVertexAttribArray {
-    uint8_t size; /* max is 4, using a bitfield might be an option */
+    unsigned normalized : 1;
+    unsigned size : 3; /* max is 4 */
     uint8_t stride;
     GLenum type;
     const void *pointer;