Question: Function return type depending on argument type

I have the following code where the useSearchResults function (it's a React hook, but it doesn't matter) initializes the state based on the argument.
config.state can be anything defined in SearchResultsState or a function that returns anything defined in SearchResultsState.
The state returned by useSearchResults is defined as Partial<SearchResultsState> because, initially could be anything in that state type. I need state to be a partial SearchResultsState but also whatever type config.state has (an object or the object returned by a function or undefined).
I've trying and trying around this but couldn't find a solution.

type State = Record<string, any>;
interface IWidget {
  State: State;
}
interface SearchResultsState extends State {
  a: number; 
  b: string;
}

interface SearchResultsWidget extends IWidget {
  State: SearchResultsState;
}
export type StateInitializerFunction<W extends IWidget> = () => Partial<W['State']>;
export type StateInitializer<W extends IWidget> = StateInitializerFunction<W> | Partial<W['State']>;

export type WidgetInitializer<W extends IWidget> = {
  state?: StateInitializer<W>;
};

export type WidgetWrapperResult<
  W extends IWidget,
> = {
  state: Partial<W['State']>;
};


const useSearchResults = (
  config: WidgetInitializer<SearchResultsWidget> = {},
): WidgetWrapperResult<SearchResultsWidget> => {
  const state = typeof config.state === 'function' ? config.state() : config.state || {};
  
  return { state };
};

const { state: { a, b } } = useSearchResults({ state: { a: 1 } });
// or const { state: { a, b } } = useSearchResults({ state: () => ({ a: 1 }) });

const test = (a: number) => console.log(a);

// error
test(a); // a is number | undefined and I'm looking to be just number because a is number in state


You can play with it here


Answer:
If you want a function's return type to depend on its argument type then you either need to overload it with multiple call signatures, or generic in some number of type parameters.  Overloads only work if you have a relatively small number of ways you want to call the function, while generics are more appropriate to represent an arbitrary relationship between input and output.

Overloads might work for this use case, if you only want to support the four possible cases for a and b being present or possibly absent, but it's tedious:
function useSearchResults(
  config: { state?: { a: number, b: string } | (() => { a: number, b: string }) }
): { state: { a: number, b: string } };
function useSearchResults(
  config: { state?: { a: number, b?: string } | (() => { a: number, b?: string }) }
): { state: { a: number, b?: string } };
function useSearchResults(
  config: { state?: { a?: number, b: string } | (() => { a?: number, b: string }) }
): { state: { a?: number, b: string } };
function useSearchResults(
  config: { state?: { a?: number, b?: string } | (() => { a?: number, b?: string }) }
): { state: { a?: number, b?: string } };

function useSearchResults(
  config: WidgetInitializer<SearchResultsWidget> = {},
): WidgetWrapperResult<SearchResultsWidget> {
  const state = typeof config.state === 'function' ? config.state() : config.state || {};

  return { state };
};

const { state: { a, b } } = useSearchResults({ state: { a: 1 } });
/* (property) state: {
  a: number;
  b?: string | undefined;
} */

useSearchResults({ state: () => ({ a: 1 }) }).state.a


That can quickly get out of hand, though, so you might need to use generics instead.  Once you use generics you need to make the generic type arguments inferrable from the inputs, but you can't do that with indexed access types like W["state"] (there was a pull request at ms/TS#20126 which would have made this possible, but it's not part of the language now.  A new PR at ms/TS#53017 might be merged eventually, but for now it's not.)
So you'll need to refactor significantly to use that type directly.  Possibly as shown here:
type StateInitializerFunction<S extends State> = () => S;
type StateInitializer<S extends State> = StateInitializerFunction<S> | S;
type WidgetInitializer<S extends State> = { state?: StateInitializer<S>; };
type WidgetWrapperResult<S extends State> = { state: S; };

const useSearchResults = <S extends Partial<SearchResultsState>>(
  config: WidgetInitializer<S> = {},
): WidgetWrapperResult<S & Partial<SearchResultsState>> => {
  const state = typeof config.state === 'function' ? config.state() : config.state || {};
  return { state } as any;
};

Here what I've done is replace mentions of W with just its state property type S. I've also replaced the Partial stuff in the definitions and moved it to the function. Now the function useSearchResults accepts a WidgetInitializer<S> for some S that extends Partial<SearchResultsState>, and it returns a WidgetWrapperResults<S & Partial<SearchResultsState>>.  That intersection just helps make sure the compiler is still aware of b's existence when you don't pass it.  Let's test it:
const { state: { a, b } } = useSearchResults({ state: { a: 1 } });
/* (property) state: {
    a: number;
} & Partial<SearchResultsState> */

const test = (a: number) => console.log(a);

useSearchResults({ state: () => ({ a: 1 }) }).state.a

test(a);

Looks good!
Playground link to code