diff --git a/lib/simple_ui/core/builder.lua b/lib/simple_ui/core/builder.lua
index 48d80d1..437cb0c 100644
--- a/lib/simple_ui/core/builder.lua
+++ b/lib/simple_ui/core/builder.lua
@@ -34,6 +34,20 @@ function builder:makeKey(element)
     return element.type .. "<" .. tostring(element.key) .. ">"
 end
 
+--- @generic T
+--- @param element StatefulElement<T>
+--- @return T
+function builder:initStateOf(element)
+    return element.initState and element:initState() or {}
+end
+
+--- @generic T
+--- @param element StatefulElement<T>
+--- @return T
+function builder:getStateOf(element)
+    return element.getState and element:getState() or self.states[self:makeKey(element)]
+end
+
 --- @param cur UIElement
 --- @private
 function builder:build_step(cur)
@@ -43,9 +57,9 @@ function builder:build_step(cur)
     local key = self:makeKey(cur)
     if key then
         cur = cur --[[@as StatefulElement]]
-        local storedState = self.states[key]
+        local storedState = self:getStateOf(cur)
         if not storedState then
-            self.states[key] = cur.state
+            self.states[key] = self:initStateOf(cur)
         else
             cur.state = storedState
         end
diff --git a/lib/simple_ui/core/stateful_element.lua b/lib/simple_ui/core/stateful_element.lua
index f848b05..c075729 100644
--- a/lib/simple_ui/core/stateful_element.lua
+++ b/lib/simple_ui/core/stateful_element.lua
@@ -2,14 +2,15 @@ local SingleChildElement = require "lib.simple_ui.core.single_child_element"
 
 --- @generic T : table
 --- @class StatefulElement<T> : SingleChildElement
---- @field state T
+--- @field initState? fun(self: StatefulElement<T>): T Создает исходное состояние элемента, когда он попадает в дерево в первый раз.
+--- @field getState? fun(self: StatefulElement<T>): T Возвращает состояние элемента. Можно переопределить, чтобы хранить его где хочется.
 local element = setmetatable({}, SingleChildElement)
 element.__index = element
 element.type = "StatefulElement"
 element.state = {}
 
 --- @return StatefulElement<T>
---- @param values {key: any, state: T?, child: UIElement?}
+--- @param values {key: any, child: UIElement?}
 function element:new(values)
     return SingleChildElement.new(self, values)
 end
diff --git a/lib/simple_ui/level/test.lua b/lib/simple_ui/level/test.lua
index 910fdc8..3df0ac6 100644
--- a/lib/simple_ui/level/test.lua
+++ b/lib/simple_ui/level/test.lua
@@ -16,8 +16,9 @@ local Canary = setmetatable({}, StatefulElement)
 Canary.__index = Canary
 Canary.type = "Canary"
 
---- @class CanaryState
-Canary.state = { i = 0 }
+function Canary:initState()
+    return { i = self.key == "canary1" and 0 or 100 }
+end
 
 function Canary:build()
     self.state.i = self.state.i and self.state.i + 1 or 0
@@ -62,9 +63,6 @@ function MyWidget:build()
                             height = 100,
                             child = Canary:new {
                                 key = "canary2",
-                                state = {
-                                    i = 100
-                                }
                             }
                         },
                     },