How to make typecript know the return type of a function when passing different variables?

I want to implement a set of 3 functions: pair, head, tail that works like this

const p = pair(1, "hello")
const h = head(p) // 1
const t = tail(p) // "hello"

I have 2 implementations:

// version 1: using array data structure
function pair<T1, T2>(a: T1, b: T2) {
  const result: [T1, T2] = [a, b]
  return result
}

function head<T1, T2>(p: [T1, T2]) {
  return p[0]
}

function tail<T1, T2>(p: [T1, T2]) {
  return p[1]
}

const p = pair(1, "hello")
const h = head(p) // typescript know h is a number
const t = tail(p) // typescript know t is a string

This version works perfectly, typescript know exactly what type of h and t

// version 2: using closures
type Pair<T1, T2> = (n: 1 | 2) => T1 | T2

function pair<T1, T2>(a: T1, b: T2) {
  const result: Pair<T1, T2> = (n: 1 | 2) => {
    switch (n) {
      case 1:
        return a
      case 2:
        return b
    }
  }
  return result
}

function head<T1, T2>(p: Pair<T1, T2>) {
  return p(1)
}

function tail<T1, T2>(p: Pair<T1, T2>) {
  return p(2)
}

const p = pair(1, 'hello')
const h = head(p) // typescript says h is number | string
const t = tail(p) // typescript says t is number | string

How can I improve the version 2 to make typescript know exactly what type of h and t?

Using overloading and extends you can achieve the desired behaviour

// version 2: using closures
// Overload Pair to match version 1 signature

interface Pair<T1, T2>{
  (n: 1): T1;
  (n: 2): T2;
  (n: 1 | 2): T1 | T2;
}


function pair<T1, T2>(a: T1, b: T2): Pair<T1, T2> {
  const result = (n: 1 | 2) => {
    switch (n) {
      case 1:
        return a ;
      case 2:
        return b ;
    }
  }
  return result as Pair<T1, T2>;
}

function head<T1, T2>(p: Pair<T1, T2>) {
  return p(1)
}

function tail<T1, T2>(p: Pair<T1, T2>) {
  return p(2)
}

const p = pair(1, 'hello')
const h = head(p) // typescript says h is number
const t = tail(p) // typescript says t is string