AllenDang · AllenDang · Aug 8, 2022 · Aug 6, 2022 · Aug 6, 2022
diff --git a/ClickableWidgets.go b/ClickableWidgets.go
@@ -329,17 +329,14 @@ func (b *ImageButtonWithRgbaWidget) FramePadding(padding int) *ImageButtonWithRg
 
 // Build implements Widget interface.
 func (b *ImageButtonWithRgbaWidget) Build() {
-	if state := Context.GetState(b.id); state == nil {
-		Context.SetState(b.id, &imageState{})
+	if state := GetState[imageState](Context, b.id); state == nil {
+		SetState(&Context, b.id, &imageState{})
 
 		NewTextureFromRgba(b.rgba, func(tex *Texture) {
-			Context.SetState(b.id, &imageState{texture: tex})
+			SetState(&Context, b.id, &imageState{texture: tex})
 		})
 	} else {
-		var isOk bool
-		imgState, isOk := state.(*imageState)
-		Assert(isOk, "ImageButtonWithRgbaWidget", "Build", "got unexpected type of widget's state")
-		b.ImageButtonWidget.texture = imgState.texture
+		b.ImageButtonWidget.texture = state.texture
 	}
 
 	b.ImageButtonWidget.Build()

diff --git a/CodeEditor.go b/CodeEditor.go
@@ -204,17 +204,12 @@ func (ce *CodeEditorWidget) Build() {
 }
 
 func (ce *CodeEditorWidget) getState() (state *codeEditorState) {
-	if s := Context.GetState(ce.title); s == nil {
+	if state = GetState[codeEditorState](Context, ce.title); state == nil {
 		state = &codeEditorState{
 			editor: imgui.NewTextEditor(),
 		}
 
-		Context.SetState(ce.title, state)
-	} else {
-		var isOk bool
-		state, isOk = s.(*codeEditorState)
-		Assert(isOk, "CodeEditorWidget", "getState", "unexpected widget's state type")
+		SetState(&Context, ce.title, state)
 	}
-
 	return state
 }
diff --git a/Context.go b/Context.go
@@ -1,6 +1,7 @@
 package giu
 
 import (
+	"fmt"
 	"sync"
 
 	"github.com/AllenDang/imgui-go"
@@ -16,6 +17,11 @@ type Disposable interface {
 	Dispose()
 }
 
+type genericDisposable[T any] interface {
+	Disposable
+	*T
+}
+
 type state struct {
 	valid bool
 	data  Disposable
@@ -101,19 +107,40 @@ func (c *context) cleanState() {
 	c.widgetIndexCounter = 0
 }
 
+func SetState[T any, PT genericDisposable[T]](c *context, id string, data PT) {
+	c.state.Store(id, &state{valid: true, data: data})
+}
+
 func (c *context) SetState(id string, data Disposable) {
 	c.state.Store(id, &state{valid: true, data: data})
 }
 
+func GetState[T any, PT genericDisposable[T]](c context, id string) PT {
+	if s, ok := c.load(id); ok {
+		s.valid = true
+		data, isOk := s.data.(PT)
+		Assert(isOk, "Context", "GetState", fmt.Sprintf("got state of unexpected type: expected %T, instead found %T", new(T), s.data))
+		return data
+
+	}
+	return nil
+}
+
 func (c *context) GetState(id string) any {
+	if s, ok := c.load(id); ok {
+		s.valid = true
+		return s.data
+	}
+	return nil
+}
+
+func (c *context) load(id any) (*state, bool) {
 	if v, ok := c.state.Load(id); ok {
 		if s, ok := v.(*state); ok {
-			s.valid = true
-			return s.data
+			return s, true
 		}
 	}
-
-	return nil
+	return nil, false
 }
 
 // Get widget index for current layout.

diff --git a/Context_test.go b/Context_test.go
@@ -12,10 +12,16 @@ func (s *teststate) Dispose() {
 	// noop
 }
 
+type teststate2 struct{}
+
+func (t *teststate2) Dispose() {
+	// noop
+}
+
 func Test_SetGetState(t *testing.T) {
 	tests := []struct {
 		id   string
-		data Disposable
+		data *teststate
 	}{
 		{"nil", nil},
 		{"pointer", &teststate{}},
@@ -24,36 +30,69 @@ func Test_SetGetState(t *testing.T) {
 	for _, tc := range tests {
 		t.Run(tc.id, func(t *testing.T) {
 			ctx := context{}
-			ctx.SetState(tc.id, tc.data)
-			restored := ctx.GetState(tc.id)
+			SetState(&ctx, tc.id, tc.data)
+			restored := GetState[teststate](ctx, tc.id)
 			assert.Equal(t, tc.data, restored, "unexpected state restored")
 		})
 	}
 }
 
+func Test_SetGetStateGeneric(t *testing.T) {
+	tests := []struct {
+		id   string
+		data *teststate
+	}{
+		{"nil", nil},
+		{"pointer", &teststate{}},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.id, func(t *testing.T) {
+			ctx := context{}
+			SetState(&ctx, tc.id, tc.data)
+			restored := GetState[teststate](ctx, tc.id)
+			assert.Equal(t, tc.data, restored, "unexpected state restored")
+		})
+	}
+}
+
+func Test_SetGetWrongStateGeneric(t *testing.T) {
+	id := "id"
+	data := &teststate{}
+	ctx := context{}
+
+	defer func() {
+		if r := recover(); r == nil {
+			t.Errorf("expected code to assert to panic, but it didn't")
+		}
+	}()
+	SetState(&ctx, id, data)
+	GetState[teststate2](ctx, id)
+}
+
 func Test_invalidState(t *testing.T) {
 	ctx := context{}
 
 	state1ID := "state1"
 	state2ID := "state2"
-	states := map[string]Disposable{
-		state1ID: &teststate{},
-		state2ID: &teststate{},
+	states := map[string]*teststate{
+		state1ID: {},
+		state2ID: {},
 	}
 
 	for i, s := range states {
-		ctx.SetState(i, s)
+		SetState(&ctx, i, s)
 	}
 
 	ctx.invalidAllState()
 
-	_ = ctx.GetState(state2ID)
+	_ = GetState[teststate](ctx, state2ID)
 
 	ctx.cleanState()
 
-	assert.NotNil(t, ctx.GetState(state2ID),
+	assert.NotNil(t, GetState[teststate](ctx, state2ID),
 		"altought state has been accessed during the frame, it has ben deleted by invalidAllState/cleanState")
-	assert.Nil(t, ctx.GetState(state1ID),
+	assert.Nil(t, GetState[teststate](ctx, state1ID),
 		"altought state hasn't been accessed during the frame, it hasn't ben deleted by invalidAllState/cleanState")
 }
 

diff --git a/EventHandler.go b/EventHandler.go
@@ -122,14 +122,9 @@ func (eh *EventHandler) Build() {
 	if eh.onActivate != nil || eh.onDeactivate != nil {
 		var state *eventHandlerState
 		stateID := GenAutoID("eventHandlerState")
-		if s := Context.GetState(stateID); s != nil {
-			var isOk bool
-			state, isOk = s.(*eventHandlerState)
-			Assert(isOk, "EventHandler", "Build", "unexpected type of state received")
-		} else {
-			newState := &eventHandlerState{}
-			Context.SetState(stateID, newState)
-			state = newState
+		if state = GetState[eventHandlerState](Context, stateID); state == nil {
+			state = &eventHandlerState{}
+			SetState(&Context, stateID, state)
 		}
 
 		if eh.onActivate != nil && isActive && !state.isActive {

diff --git a/ExtraWidgets.go b/ExtraWidgets.go
@@ -435,13 +435,9 @@ func (l *ListBoxWidget) OnMenu(onMenu func(selectedIndex int, menu string)) *Lis
 // nolint:gocognit // will fix later
 func (l *ListBoxWidget) Build() {
 	var state *ListBoxState
-	if s := Context.GetState(l.id); s == nil {
+	if state = GetState[ListBoxState](Context, l.id); state == nil {
 		state = &ListBoxState{selectedIndex: 0}
-		Context.SetState(l.id, state)
-	} else {
-		var isOk bool
-		state, isOk = s.(*ListBoxState)
-		Assert(isOk, "ListBoxWidget", "Build", "wrong state type recovered")
+		SetState(&Context, l.id, state)
 	}
 
 	child := Child().Border(l.border).Size(l.width, l.height).Layout(Layout{

diff --git a/ImageWidgets.go b/ImageWidgets.go
@@ -160,17 +160,13 @@ func (i *ImageWithRgbaWidget) OnClick(cb func()) *ImageWithRgbaWidget {
 func (i *ImageWithRgbaWidget) Build() {
 	if i.rgba != nil {
 		var imgState *imageState
-		if state := Context.GetState(i.id); state == nil {
+		if imgState = GetState[imageState](Context, i.id); imgState == nil {
 			imgState = &imageState{}
-			Context.SetState(i.id, imgState)
+			SetState(&Context, i.id, imgState)
 
 			NewTextureFromRgba(i.rgba, func(tex *Texture) {
 				imgState.texture = tex
 			})
-		} else {
-			var isOk bool
-			imgState, isOk = state.(*imageState)
-			Assert(isOk, "ImageWithRgbaWidget", "Build", "unexpected type of widget's state recovered")
 		}
 
 		i.img.texture = imgState.texture
@@ -221,21 +217,18 @@ func (i *ImageWithFileWidget) OnClick(cb func()) *ImageWithFileWidget {
 
 // Build implements Widget interface.
 func (i *ImageWithFileWidget) Build() {
-	imgState := &imageState{}
-	if state := Context.GetState(i.id); state == nil {
+	var imgState *imageState
+	if imgState = GetState[imageState](Context, i.id); imgState == nil {
 		// Prevent multiple invocation to LoadImage.
-		Context.SetState(i.id, imgState)
+		imgState = &imageState{}
+		SetState(&Context, i.id, imgState)
 
 		img, err := LoadImage(i.imgPath)
 		if err == nil {
 			NewTextureFromRgba(img, func(tex *Texture) {
 				imgState.texture = tex
 			})
 		}
-	} else {
-		var isOk bool
-		imgState, isOk = state.(*imageState)
-		Assert(isOk, "ImageWithFileWidget", "Build", "wrong type of widget's state got")
 	}
 
 	i.img.texture = imgState.texture
@@ -313,17 +306,17 @@ func (i *ImageWithURLWidget) LayoutForFailure(widgets ...Widget) *ImageWithURLWi
 
 // Build implements Widget interface.
 func (i *ImageWithURLWidget) Build() {
-	imgState := &imageState{}
-
-	if state := Context.GetState(i.id); state == nil {
-		Context.SetState(i.id, imgState)
+	var imgState *imageState
+	if imgState = GetState[imageState](Context, i.id); imgState == nil {
+		imgState = &imageState{}
+		SetState(&Context, i.id, imgState)
 
 		// Prevent multiple invocation to download image.
 		downloadContext, cancalFunc := ctx.WithCancel(ctx.Background())
-		Context.SetState(i.id, &imageState{loading: true, cancel: cancalFunc})
+		SetState(&Context, i.id, &imageState{loading: true, cancel: cancalFunc})
 
 		errorFn := func(err error) {
-			Context.SetState(i.id, &imageState{failure: true})
+			SetState(&Context, i.id, &imageState{failure: true})
 
 			// Trigger onFailure event
 			if i.onFailure != nil {
@@ -361,7 +354,7 @@ func (i *ImageWithURLWidget) Build() {
 			rgba := ImageToRgba(img)
 
 			NewTextureFromRgba(rgba, func(tex *Texture) {
-				Context.SetState(i.id, &imageState{
+				SetState(&Context, i.id, &imageState{
 					loading: false,
 					failure: false,
 					texture: tex,
@@ -373,10 +366,6 @@ func (i *ImageWithURLWidget) Build() {
 				i.onReady()
 			}
 		}()
-	} else {
-		var isOk bool
-		imgState, isOk = state.(*imageState)
-		Assert(isOk, "ImageWithURLWidget", "Build", "wrong type of widget's state recovered.")
 	}
 
 	switch {

diff --git a/MemoryEditor.go b/MemoryEditor.go
@@ -44,16 +44,12 @@ func (me *MemoryEditorWidget) Build() {
 }
 
 func (me *MemoryEditorWidget) getState() (state *memoryEditorState) {
-	if s := Context.GetState(me.id); s == nil {
+	if state = GetState[memoryEditorState](Context, me.id); state == nil {
 		state = &memoryEditorState{
 			editor: imgui.NewMemoryEditor(),
 		}
 
-		Context.SetState(me.id, state)
-	} else {
-		var ok bool
-		state, ok = s.(*memoryEditorState)
-		Assert(ok, "MemoryEditorWidget", "getState", "incorrect state type recovered.")
+		SetState(&Context, me.id, state)
 	}
 
 	return state

diff --git a/Msgbox.go b/Msgbox.go
@@ -103,15 +103,9 @@ func PrepareMsgbox() Layout {
 			var state *msgboxState
 
 			// Register state.
-			stateRaw := Context.GetState(msgboxID)
-
-			if stateRaw == nil {
+			if state = GetState[msgboxState](Context, msgboxID); state == nil {
 				state = &msgboxState{title: "Info", content: "Content", buttons: MsgboxButtonsOk, resultCallback: nil, open: false}
-				Context.SetState(msgboxID, state)
-			} else {
-				var isOk bool
-				state, isOk = stateRaw.(*msgboxState)
-				Assert(isOk, "MsgboxWidget", "PrepareMsgbox", "got state of unexpected type")
+				SetState(&Context, msgboxID, state)
 			}
 
 			if state.open {
@@ -122,7 +116,7 @@ func PrepareMsgbox() Layout {
 			PopupModal(fmt.Sprintf("%s%s", state.title, msgboxID)).Layout(
 				Custom(func() {
 					// Ensure the state is valid.
-					Context.GetState(msgboxID)
+					GetState[msgboxState](Context, msgboxID)
 				}),
 				Label(state.content).Wrapped(true),
 				buildMsgboxButtons(state.buttons, state.resultCallback),
@@ -135,15 +129,12 @@ func PrepareMsgbox() Layout {
 type MsgboxWidget struct{}
 
 func (m *MsgboxWidget) getState() *msgboxState {
-	stateRaw := Context.GetState(msgboxID)
-	if stateRaw == nil {
+	state := GetState[msgboxState](Context, msgboxID)
+	if state == nil {
 		panic("Msgbox is not prepared. Invoke giu.PrepareMsgbox in the end of the layout.")
 	}
 
-	result, isOk := stateRaw.(*msgboxState)
-	Assert(isOk, "MsgboxWidget", "getState", "unexpected type of widget's state recovered")
-
-	return result
+	return state
 }
 
 // Msgbox opens message box.