Сделать функцию резолвер для размеров

нужно сделать функцию резолвер для размеров по аналогии с резолвером цветов

import React, { cloneElement, ReactElement, useMemo } from 'react';
import { ClassNameProps, component, overridable, useTheme } from '@spatium/styled-system';
import { Mode, ModifiersModifierProps, PaddingModifierProps, RadiusModifierProps, RowModifierProps, Theme, theme, ThemeColor, ThemeSpace, IconProps } from '@spatium/theme';
import type { ButtonPosition, ButtonState } from '../lib/button';
import { RectangleButton } from '../rectangle-button';

export const buttonSizeList = ['large', 'medium', 'small'] as const;
type ButtonSize = typeof buttonSizeList[number];

export const buttonVariantList = ['primary', 'secondary'] as const;
export type SquareButtonVariant = typeof buttonVariantList[number];

export const typeList = ['square', 'round'] as const;
export type SquareButtonType = typeof typeList[number];

export type SquareButtonProps = {
  /**
 * Child element of Square Button
 * @example <SquareButton ... children={<SomeReactElement />} />
 */
  children: ReactElement<IconProps>;
  /**
 * Sets state of Square Button
 *
 * @default state = 'normal'
 * @example <SquareButton ... state={isLoading ? 'loading' : 'normal'} />
 */
  state?: ButtonState;
  /**
 * Sets one of two positions of components: on background
 * or on a surface
 * @default position = 'background'
 * @example <SquareButton ... position='background' />
 */
  position?: ButtonPosition;
  /**
 * Sets one of the default sizes of Square Button
 * @default size = 'large'
 * @example <SquareButton ... size='large' />
 */
  size?: ButtonSize;
  /**
 * Sets one of two default variants of buttons: primary and secondary
 * @default variant = 'primary'
 * @example <SquareButton ... variant={isActive ? 'primary' : 'secondary'} />
 */
  variant?: SquareButtonVariant;
  /**
 * Sets one of two default types of Square Button: round and square
 * @default type = 'round'
 * @example <SquareButton ... children={<SomeReactElement />} />
 */
  type?: SquareButtonType;
  /**
 * Allows you to customize Square Button size
 *
 * @example <SquareButton ... themeSize={theme.space(100)} />
 */
  themeSize?: ThemeSpace;
  /**
 * Allows you to customize Square Button background color
 *
 * @example <SquareButton ... themeBackgroundColor={theme.colors.red} />
 */
  themeBackgroundColor?: ThemeColor;
  /**
 * Allows you to customize Square Button content color
 *
 * @example <SquareButton ... themeContentColor={theme.colors.blue} />
 */
  themeContentColor?: ThemeColor;
} & ModifiersModifierProps & RowModifierProps & PaddingModifierProps & RadiusModifierProps & ClassNameProps;

export const SquareButton = component<SquareButtonProps>(({
  state = 'normal',
  position = 'background',
  size = 'large',
  variant = 'primary',
  children,
  themeSize = sizeMap[size],
  type = 'round',
  themeRadius = type === 'square' ? theme.space(10) : themeSize.divided(2),
  themePadding,
  themeBackgroundColor,
  themeContentColor,
  ...rest
}) => {
  const { mode } = useTheme<Theme>();

  const { backgroundColor, contentColor } = useMemo(() => {
    const { background, content } = resolveSquareButtonColor(state, variant, mode, position);
    return {
      backgroundColor: themeBackgroundColor ?? background ?? undefined,
      contentColor: themeContentColor ?? content ?? undefined,
    };
  }, [themeBackgroundColor, state, variant, mode, position, themeContentColor]);

  const icon = useMemo(() => (
    cloneElement(children, {
      ...children.props,
      themeColor: contentColor,
      themeSize: iconSizeMap[size],
    })
  ), [contentColor, children, size]);

  return (
    <RectangleButton
      state={state}
      position={position}
      variant={variant}
      themeWidth={themeSize}
      themeHeight={themeSize}
      themeRadius={themeRadius}
      themeBackgroundColor={backgroundColor}
      themePadding={themePadding}
      {...rest}
    >
      {icon}
    </RectangleButton>
  );
});

const disabledBackground: Record<Mode, Node> = {
  light: theme.colors.surface.alt(1),
  dark: {
    background: theme.colors.surface.alt(1),
    surface: theme.colors.surface,
  },
};

const defaultBackgroundMap: Record<ButtonState, Node> = {
  normal: {
    primary: {
      light: theme.colors.primary,
      dark: {
        background: theme.colors.surface,
        surface: theme.colors.surface.alt(2),
      },
    },
    secondary: {
      light: theme.colors.surface,
      dark: {
        background: theme.colors.surface,
        surface: theme.colors.surface.alt(2),
      },
    },
  },
  pressed: {
    primary: {
      light: theme.colors.primary.alt(1),
      dark: {
        background: theme.colors.surface.alt(1),
        surface: theme.colors.surface,
      },
    },
    secondary: {
      light: theme.colors.surface.alt(2),
      dark: {
        background: theme.colors.surface.alt(1),
        surface: theme.colors.surface,
      },
    },
  },
  disabled: {
    primary: disabledBackground,
    secondary: disabledBackground,
  },
  active: {
    primary: theme.colors.surface.alt(3),
    secondary: theme.colors.surface.alt(3),
  },
  loading: {
    primary: disabledBackground,
    secondary: disabledBackground,
  },
  success: {
    primary: disabledBackground,
    secondary: disabledBackground,
  },
};

const defaultContentMap: Record<ButtonState, Node> = {
  normal: {
    light: theme.colors.foreground,
    dark: theme.colors.foreground,
  },
  pressed: {
    light: theme.colors.foreground,
    dark: theme.colors.foreground,
  },
  disabled: {
    light: theme.colors.foreground,
    dark: theme.colors.foreground,
  },
  active: {
    primary: theme.colors.foreground.alt(3),
    secondary: theme.colors.foreground.alt(3),
  },
  loading: theme.colors.foreground,
  success: theme.colors.success,
};

type StateNode = { [key in ButtonState]: Node };
type VariantNode = { [key in SquareButtonVariant]: Node };
type ModeNode = { [key in Mode]: Node };
type PositionNode = { [key in ButtonPosition]: Node };
type Node = null | ThemeColor | StateNode | VariantNode | ModeNode | PositionNode;

export const resolveSquareButtonColor = overridable((state: ButtonState, variant: SquareButtonVariant, mode: Mode, position: ButtonPosition) => {
  const resolveNode = (node: Node): ThemeColor | null => {
    if (node && 'normal' in node) {
      return resolveNode(node[state]);
    } else if (node && 'primary' in node) {
      return resolveNode(node[variant]);
    } else if (node && 'light' in node) {
      return resolveNode(node[mode]);
    } else if (node && 'background' in node) {
      return resolveNode(node[position]);
    } else {
      return node;
    }
  };
  return { background: resolveNode(defaultBackgroundMap), content: resolveNode(defaultContentMap) };
});

const sizeMap: Record<ButtonSize, ThemeSpace> = {
  large: theme.space(60),
  medium: theme.space(40),
  small: theme.space(36),
};

const iconSizeMap: Record<ButtonSize, ThemeSpace> = {
  large: theme.space(36),
  medium: theme.space(24),
  small: theme.space(24),
};