Merge pull request #13 from mdgriffith/twop/func-unwrapping

Twop/func unwrapping

src/compile-testcases.ts

@@ -23,6 +23,7 @@ import {
   convertFunctionExpressionsToArrowFuncs,
 } from './experiments/modernizeJS';
 import { createRemoveUnusedLocalsTransform } from './experiments/removeUnusedLocals';
+import { createPassUnwrappedFunctionsTransformer } from './experiments/passUnwrappedFunctions';
 
 export const compileAndTransform = async (
   dir: string,
@@ -64,11 +65,15 @@ export const compileAndTransform = async (
     Mode.Prod
   );
 
+  let inlineCtx: InlineContext | undefined;
   const transformations = removeDisabled([
     [options.variantShapes, normalizeVariantShapes],
     [
       options.inlineFunctions,
-      createFunctionInlineTransformer(reportInlineTransformResult),
+      createFunctionInlineTransformer(ctx => {
+        inlineCtx = ctx;
+        reportInlineTransformResult(ctx);
+      }),
     ],
     [options.inlineEquality, inlineEquality()],
     [
@@ -77,6 +82,10 @@ export const compileAndTransform = async (
         InlineMode.UsingLiteralObjects(Mode.Prod)
       ),
     ],
+    [
+      options.passUnwrappedFunctions,
+      createPassUnwrappedFunctionsTransformer(() => inlineCtx),
+    ],
     includeObjectUpdate(options.objectUpdate),
     [options.arrowFns, convertFunctionExpressionsToArrowFuncs],
     [options.unusedValues, createRemoveUnusedLocalsTransform()],
@@ -207,5 +216,6 @@ function reportInlineTransformResult(ctx: InlineContext) {
   console.log(
     `functionInlineTransformer: splitCount=${splits.size}, partialApplicationCount=${partialApplications.size}`,
     inlined
+    // splits
   );
 }

src/experiments/createTSprogram.ts

@@ -0,0 +1,63 @@
+import ts from 'typescript';
+
+const cache = new Map<string, ts.SourceFile | undefined>();
+const defaultCompilerHost = ts.createCompilerHost({});
+
+function serveLibFile(
+  name: string,
+  languageVersion: ts.ScriptTarget
+): ts.SourceFile | undefined {
+  const cached = cache.get(name);
+  if (cached) return cached;
+
+  const val = defaultCompilerHost.getSourceFile(name, languageVersion);
+  cache.set(name, val);
+  return val;
+}
+
+const printer = ts.createPrinter();
+
+export const createProgramFromSource = (
+  source: ts.SourceFile
+): [ts.Program, ts.SourceFile] => {
+  const sourceCopy = ts.createSourceFile(
+    source.fileName,
+    printer.printFile(source),
+    ts.ScriptTarget.ES2018
+  );
+
+  const customCompilerHost: ts.CompilerHost = {
+    getSourceFile: (name, languageVersion) => {
+      // console.log(`getSourceFile ${name}`);
+
+      if (name === sourceCopy.fileName) {
+        return sourceCopy;
+      } else {
+        return serveLibFile(name, languageVersion);
+      }
+    },
+    writeFile: () => {},
+    getDefaultLibFileName: () =>
+      'node_modules/typescript/lib/lib.es2018.full.d.ts',
+    useCaseSensitiveFileNames: () => false,
+    getCanonicalFileName: filename => filename,
+    getCurrentDirectory: () => '',
+    getNewLine: () => '\n',
+    getDirectories: () => [],
+    fileExists: () => true,
+    readFile: () => '',
+  };
+
+  const program = ts.createProgram(
+    [sourceCopy.fileName],
+    {
+      allowJs: true,
+      noUnusedLocals: true,
+      checkJs: true,
+      outDir: 'yo',
+    },
+    customCompilerHost
+  );
+
+  return [program, sourceCopy];
+};

src/experiments/inlineWrappedFunctions.ts

@@ -1,4 +1,5 @@
 import ts from 'typescript';
+// import { matchElmSource } from './patterns';
 
 /*
 
@@ -75,11 +76,7 @@ function reportInlinining(split: FuncSplit, { inlined }: InlineContext) {
   }
 }
 
-export const createFunctionInlineTransformer = (
-  reportResult?: (res: InlineContext) => void
-): ts.TransformerFactory<ts.SourceFile> => context => {
-  return sourceFile => {
-    const inlineContext: InlineContext = {
+export const createInlineContext = (): InlineContext => ({
   functionsThatWrapFunctions: new Map(),
   splits: new Map(),
   partialApplications: new Map(),
@@ -89,8 +86,16 @@ export const createFunctionInlineTransformer = (
     fromWrapper: 0,
     partialApplications: 0,
   },
-    };
+});
 
+export const createFunctionInlineTransformer = (
+  reportResult?: (res: InlineContext) => void
+): ts.TransformerFactory<ts.SourceFile> => context => {
+  return sourceFile => {
+    const inlineContext: InlineContext = createInlineContext();
+
+    // todo hack to only inline top level functions
+    // const { topScope } = matchElmSource(sourceFile)!;
     const splitter = createSplitterVisitor(inlineContext, context);
     const splittedNode = ts.visitNode(sourceFile, splitter);
 
@@ -106,16 +111,33 @@ export const createFunctionInlineTransformer = (
   };
 };
 
+const isTopLevelScope = (path: ts.Node[]) => {
+  const funcExpCount = path.reduce(
+    (c, n) => (ts.isFunctionExpression(n) ? c + 1 : c),
+    0
+  );
+  const funcDeclCount = path.reduce(
+    (c, n) => (ts.isFunctionDeclaration(n) ? c + 1 : c),
+    0
+  );
+  // meaning top level scope function body
+  return funcExpCount === 1 && funcDeclCount === 0;
+};
+
 const createSplitterVisitor = (
   { splits, partialApplications, functionsThatWrapFunctions }: InlineContext,
   context: ts.TransformationContext
 ) => {
-  const visitor = (node: ts.Node): ts.VisitResult<ts.Node> => {
+  const visitor = (path: ts.Node[]) => (
+    node: ts.Node
+  ): ts.VisitResult<ts.Node> => {
     // detects "var a"
     if (
       ts.isVariableDeclaration(node) &&
+      // todo this is basically a hack to only be able to inline top level functions
       ts.isIdentifier(node.name) &&
-      node.initializer
+      node.initializer &&
+      isTopLevelScope(path)
     ) {
       // detects an alias to existing split
       if (ts.isIdentifier(node.initializer)) {
@@ -176,13 +198,12 @@ const createSplitterVisitor = (
                 }
               }
 
-              // and it is a function
-              // detects "var a = F123( function (a) {return a})"
-              // or "var a = F123( a => a)"
-              if (
-                ts.isArrowFunction(maybeFuncExpression) ||
-                ts.isFunctionExpression(maybeFuncExpression)
-              ) {
+              // it can be either a function expression:
+              // "var a = F123( function (a) {return a})"
+              // "var a = F123( a => a)"
+              // something like
+              // var a = F2(Math.pow)
+              if (true) {
                 // TODO typecheck?
                 const rawLambdaName = deriveRawLambdaName(originalName);
 
@@ -207,12 +228,21 @@ const createSplitterVisitor = (
                   ])
                 );
 
+                // if it is a function expression check if we need an unwrapped version of it
+                if (
+                  ts.isArrowFunction(maybeFuncExpression) ||
+                  ts.isFunctionExpression(maybeFuncExpression)
+                ) {
                   const maybeWrapper = checkIfFunctionReturnsWrappedFunction(
                     maybeFuncExpression
                   );
 
                   if (maybeWrapper) {
-                  functionsThatWrapFunctions.set(node.name.text, maybeWrapper);
+                    functionsThatWrapFunctions.set(
+                      node.name.text,
+                      maybeWrapper
+                    );
+                  }
                 }
 
                 return [lambdaDeclaration, newDeclaration];
@@ -341,10 +371,10 @@ const createSplitterVisitor = (
       }
     }
 
-    return ts.visitEachChild(node, visitor, context);
+    return ts.visitEachChild(node, visitor(path.concat(node)), context);
   };
 
-  return visitor;
+  return visitor([]);
 };
 
 const createInlinerVisitor = (

src/experiments/modernizeJS.ts

@@ -22,7 +22,7 @@ const _Utils_update = (oldRecord, updatedFields) => (Object.assign({}, oldRecord
 }
 `;
 
-const extractBody = (sourceText: string): ts.Node => {
+export const extractAstFromCode = (sourceText: string): ts.Node => {
   const source = ts.createSourceFile('bla', sourceText, ts.ScriptTarget.ES2018);
   return source.statements[0];
 };
@@ -39,11 +39,11 @@ export const createReplaceUtilsUpdateWithObjectSpread = (
       ) {
         switch (kind) {
           case ObjectUpdate.UseSpreadForUpdateAndOriginalRecord:
-            return extractBody(spreadForBoth);
+            return extractAstFromCode(spreadForBoth);
           case ObjectUpdate.UseSpreadOnlyToMakeACopy:
-            return extractBody(copyWithSpread);
+            return extractAstFromCode(copyWithSpread);
           case ObjectUpdate.UseAssign:
-            return extractBody(assign);
+            return extractAstFromCode(assign);
         }
       }
 

src/experiments/passUnwrappedFunctions.ts

@@ -0,0 +1,264 @@
+import ts from 'typescript';
+// import { createFunctionInlineTransformer } from './inlineWrappedFunctions';
+import { matchWrappedInvocation, matchWrapping } from './patterns';
+import { createProgramFromSource } from './createTSprogram';
+import { InlineContext } from './inlineWrappedFunctions';
+
+const deriveNewFuncName = (funcName: string) => funcName + '_unwrapped';
+
+type WrappedUsage = {
+  arity: number; // expected arity of the function being unwrapped inside
+  funcDeclaration: ts.VariableDeclaration;
+  funcName: string;
+  parameterName: string;
+  parameterPos: number;
+  funcExpression: ts.FunctionExpression;
+  callExpression: ts.CallExpression;
+};
+
+export const createPassUnwrappedFunctionsTransformer = (
+  getCtx: () => InlineContext | undefined
+): ts.TransformerFactory<ts.SourceFile> => context => {
+  getCtx;
+  return sourceFile => {
+    const foundFunctions: WrappedUsage[] = [];
+    const [program, copiedSource] = createProgramFromSource(sourceFile);
+    const typeChecker = program.getTypeChecker();
+
+    const collectFunctions = (node: ts.Node): ts.VisitResult<ts.Node> => {
+      const invocation = matchWrappedInvocation(node);
+
+      if (invocation) {
+        const { arity, callExpression, calleeName: funcName } = invocation;
+
+        const symbol = typeChecker.getSymbolAtLocation(funcName);
+
+        const [declaration] = symbol?.declarations || [];
+
+        if (
+          declaration &&
+          ts.isParameter(declaration) &&
+          ts.isIdentifier(declaration.name) &&
+          ts.isFunctionExpression(declaration.parent) &&
+          ts.isVariableDeclaration(declaration.parent.parent) &&
+          ts.isIdentifier(declaration.parent.parent.name)
+        ) {
+          const funcDeclaration = declaration.parent.parent;
+          const func = declaration.parent;
+          const parameterPos = func.parameters.findIndex(
+            p => p === declaration
+          );
+          const funcName = declaration.parent.parent.name.text;
+
+          foundFunctions.push({
+            arity,
+            callExpression,
+            funcDeclaration: funcDeclaration,
+            funcExpression: func,
+            parameterPos,
+            funcName,
+            parameterName: declaration.name.text,
+          });
+        }
+      }
+
+      return ts.visitEachChild(node, collectFunctions, context);
+    };
+
+    const addFunctionsWithoutUnwrapping = (
+      node: ts.Node
+    ): ts.VisitResult<ts.Node> => {
+      if (ts.isVariableDeclaration(node)) {
+        // todo make it a map maybe?
+        const funcToModify = foundFunctions.find(
+          f => f.funcDeclaration == node
+        );
+
+        if (funcToModify) {
+          const modifyFunction = (node: ts.Node): ts.VisitResult<ts.Node> => {
+            if (ts.isCallExpression(node) && ts.isIdentifier(node.expression)) {
+              const calledFuncIdentifier = node.expression;
+              const match = matchWrappedInvocation(node);
+              if (
+                match &&
+                match.calleeName.text === funcToModify.parameterName &&
+                match.arity === funcToModify.arity
+              ) {
+                // transforms A2(func, a,b) into func(a,b)
+                return ts.createCall(
+                  match.calleeName,
+                  undefined,
+                  // recursively transform all calls within a function
+                  match.args.map(arg => ts.visitNode(arg, modifyFunction))
+                );
+              }
+
+              // recursive call to itself
+              if (calledFuncIdentifier.text === funcToModify.funcName) {
+                return ts.createCall(
+                  ts.createIdentifier(deriveNewFuncName(funcToModify.funcName)),
+                  undefined,
+                  node.arguments.map(arg => ts.visitNode(arg, modifyFunction))
+                );
+              }
+
+              // now it can be a call to another function but with passed in parameter
+              // that is now unwrapped
+              // thus check if we have an unwrapped version
+              if (
+                node.arguments.some(
+                  arg =>
+                    ts.isIdentifier(arg) &&
+                    arg.text === funcToModify.parameterName
+                )
+              ) {
+                const existingUnwrappedFunc = foundFunctions.find(
+                  f => f.funcName === calledFuncIdentifier.text
+                );
+
+                if (!existingUnwrappedFunc) {
+                  // todo we need to bail out but let's fail for now
+                  throw new Error(
+                    `${calledFuncIdentifier.text} doesn't have an unwrapped version`
+                  );
+                }
+
+                // we need to make sure that arity matches
+                if (existingUnwrappedFunc.arity !== funcToModify.arity) {
+                  // todo we need to bail out but let's fail for now
+                  throw new Error(
+                    `${calledFuncIdentifier.text} doesn't match expected arity of a passed function with expected arity of the parent function`
+                  );
+                }
+
+                return ts.createCall(
+                  ts.createIdentifier(
+                    deriveNewFuncName(calledFuncIdentifier.text)
+                  ),
+                  undefined,
+                  node.arguments.map(arg => ts.visitNode(arg, modifyFunction))
+                );
+              }
+            }
+
+            return ts.visitEachChild(node, modifyFunction, context);
+          };
+
+          const newFuncExpression = ts.visitNode(
+            funcToModify.funcExpression,
+            modifyFunction
+          );
+
+          return [
+            node,
+            ts.createVariableDeclaration(
+              deriveNewFuncName(funcToModify.funcName),
+              undefined,
+              newFuncExpression
+            ),
+          ];
+        }
+      }
+
+      return ts.visitEachChild(node, addFunctionsWithoutUnwrapping, context);
+    };
+
+    const replaceUsagesWithUnwrappedVersion = (
+      node: ts.Node
+    ): ts.VisitResult<ts.Node> => {
+      if (ts.isCallExpression(node)) {
+        const { expression } = node;
+        if (ts.isIdentifier(expression)) {
+          // todo make it a map maybe?
+          const funcToUnwrap = foundFunctions.find(
+            f => f.funcName == expression.text
+          );
+
+          if (
+            funcToUnwrap &&
+            // actually the same symbol
+            typeChecker.getSymbolAtLocation(expression) ===
+              typeChecker.getSymbolAtLocation(funcToUnwrap.funcDeclaration.name)
+          ) {
+            const args = node.arguments;
+            const argPos = funcToUnwrap.parameterPos;
+            const funcParameter: ts.Expression | undefined = args[argPos];
+
+            if (funcParameter) {
+              const match = matchWrapping(funcParameter);
+
+              // it means that it is something like (..., F3(function (a,b,c) {...}), ...)
+              if (match) {
+                return ts.createCall(
+                  ts.createIdentifier(deriveNewFuncName(expression.text)),
+                  undefined,
+                  [
+                    ...args
+                      .slice(0, funcToUnwrap.parameterPos)
+                      .map(a =>
+                        ts.visitNode(a, replaceUsagesWithUnwrappedVersion)
+                      ),
+                    ts.visitNode(
+                      match.wrappedExpression,
+                      replaceUsagesWithUnwrappedVersion
+                    ),
+                    ...args
+                      .slice(argPos + 1)
+                      .map(a =>
+                        ts.visitNode(a, replaceUsagesWithUnwrappedVersion)
+                      ),
+                  ]
+                );
+              } else if (ts.isIdentifier(funcParameter)) {
+                const existingSplit = getCtx()?.splits.get(funcParameter.text);
+                // if (funcParameter.text)
+                if (
+                  existingSplit &&
+                  existingSplit.arity === funcToUnwrap.arity
+                ) {
+                  return ts.createCall(
+                    ts.createIdentifier(deriveNewFuncName(expression.text)),
+                    undefined,
+                    [
+                      ...args
+                        .slice(0, funcToUnwrap.parameterPos)
+                        .map(a =>
+                          ts.visitNode(a, replaceUsagesWithUnwrappedVersion)
+                        ),
+                      ts.createIdentifier(existingSplit.rawLambdaName),
+                      ...args
+                        .slice(argPos + 1)
+                        .map(a =>
+                          ts.visitNode(a, replaceUsagesWithUnwrappedVersion)
+                        ),
+                    ]
+                  );
+                }
+              }
+            }
+          }
+        }
+      }
+
+      return ts.visitEachChild(
+        node,
+        replaceUsagesWithUnwrappedVersion,
+        context
+      );
+    };
+
+    ts.visitNode(copiedSource, collectFunctions);
+
+    const withUnwrappedFunctions = ts.visitNode(
+      copiedSource,
+      addFunctionsWithoutUnwrapping
+    );
+
+    const withInlinedCalls = ts.visitNode(
+      withUnwrappedFunctions,
+      replaceUsagesWithUnwrappedVersion
+    );
+
+    return withInlinedCalls;
+  };
+};

src/experiments/patterns.ts

@@ -0,0 +1,94 @@
+import ts from 'typescript';
+export type Pattern<T> = (node: ts.Node) => T | undefined;
+
+const invocationRegex = /A(?<arity>[1-9]+[0-9]*)/;
+
+const wrapperRegex = /F(?<arity>[1-9]+[0-9]*)/;
+
+type WrappedInvocation = {
+  args: ts.Expression[];
+  calleeName: ts.Identifier;
+  callExpression: ts.CallExpression;
+  arity: number;
+};
+
+export const matchWrappedInvocation: Pattern<WrappedInvocation> = node => {
+  if (ts.isCallExpression(node)) {
+    const expression = node.expression;
+    // detects f(..)
+    if (ts.isIdentifier(expression)) {
+      const maybeMatch = expression.text.match(invocationRegex);
+      // detects A123(...)
+      if (maybeMatch && maybeMatch.groups) {
+        const arity = Number(maybeMatch.groups.arity);
+
+        const allArgs = node.arguments;
+        const [funcName, ...args] = allArgs;
+
+        // detects A123(funcName, ...args)
+        if (ts.isIdentifier(funcName)) {
+          if (args.length !== arity) {
+            throw new Error(
+              `something went wrong, expected number of arguments=${arity} but got ${args.length} for ${funcName.text}`
+            );
+          }
+
+          return {
+            args,
+            calleeName: funcName,
+            callExpression: node,
+            arity,
+          };
+        }
+      }
+    }
+  }
+  return undefined;
+};
+
+type Wrapping = {
+  wrappedExpression: ts.Expression;
+  arity: number;
+};
+
+export const matchWrapping: Pattern<Wrapping> = node => {
+  if (ts.isCallExpression(node) && ts.isIdentifier(node.expression)) {
+    const maybeMatch = node.expression.text.match(wrapperRegex);
+
+    if (maybeMatch && maybeMatch.groups) {
+      return {
+        arity: Number(maybeMatch.groups.arity),
+        wrappedExpression: node.arguments[0],
+      };
+    }
+  }
+
+  return undefined;
+};
+
+// SourceFile;
+//  ExpressionStatement;
+//    ParenthesizedExpression;
+//      CallExpression;
+//        FunctionExpression;
+//          Parameter;
+//            Identifier;
+//          Block; <==== that what we are looking
+//    ThisKeyword;
+
+export const matchElmSource = (
+  source: ts.SourceFile
+): { topScope: ts.Block } | undefined => {
+  const expStatement = source.statements[0];
+  console.log(expStatement);
+  if (
+    ts.isExpressionStatement(expStatement) &&
+    ts.isParenthesizedExpression(expStatement.expression) &&
+    ts.isCallExpression(expStatement.expression.expression) &&
+    ts.isFunctionExpression(expStatement.expression.expression.expression)
+  ) {
+    return { topScope: expStatement.expression.expression.expression.body };
+  }
+
+  return undefined;
+};

src/generate-report.ts

@@ -5,30 +5,32 @@ Compiles all the test cases and runs them via webdriver to summarize the results
 
 */
 
-import { ObjectUpdate } from './types';
+import { ObjectUpdate, Transforms } from './types';
 import * as Reporting from './reporting';
 
-const defaultOptions = {
+const defaultOptions: Transforms = {
   prepack: true,
   variantShapes: true,
   inlineEquality: true,
   inlineFunctions: true,
   listLiterals: true,
+  passUnwrappedFunctions: true,
   arrowFns: true,
   objectUpdate: null,
-  unusedValues: false,
+  unusedValues: true,
 };
 
-const inlineEquality = {
-  prepack: false,
-  variantShapes: false,
-  inlineEquality: true,
-  inlineFunctions: false,
-  listLiterals: false,
-  arrowFns: false,
-  objectUpdate: null,
-  unusedValues: false,
-};
+// const defaultOptions: Transforms = {
+//   prepack: false,
+//   variantShapes: false,
+//   inlineEquality: false,
+//   inlineFunctions: true,
+//   passUnwrappedFunctions: true,
+//   listLiterals: false,
+//   arrowFns: false,
+//   objectUpdate: null,
+//   unusedValues: false,
+// };
 async function go() {
   const report = await Reporting.run([
     // Use `runWithBreakdown` if you want the breakdown

src/reporting.ts

@@ -206,6 +206,7 @@ const emptyOpts: Transforms = {
   inlineFunctions: false,
   inlineEquality: false,
   listLiterals: false,
+  passUnwrappedFunctions: false,
   arrowFns: false,
   objectUpdate: null,
   unusedValues: false,

src/types.ts

@@ -17,6 +17,7 @@ export type Transforms = {
   variantShapes: boolean;
   inlineEquality: boolean;
   inlineFunctions: boolean;
+  passUnwrappedFunctions: boolean;
   listLiterals: boolean;
   arrowFns: boolean;
   objectUpdate: ObjectUpdate | null;

test/blah.test.ts

@@ -1,38 +0,0 @@
-import ts from 'typescript';
-
-import {
-  createSplitFunctionDeclarationsTransformer,
-  FuncSplit,
-} from '../src/experiments/inlineWrappedFunctions';
-
-const elmOutput = `
-
-var $author$project$Main$Three = F3(
-  function (a, b, c) {
-      return {$: 'Three', a: a, b: b, c: c};
-  });
-`;
-
-// this is for debugging only for now
-// not a real test
-test('blah', () => {
-  const source = ts.createSourceFile(
-    'elm.js',
-    elmOutput,
-    ts.ScriptTarget.ES2018
-  );
-
-  const printer = ts.createPrinter();
-
-  const collectedSplits: FuncSplit[] = [];
-  const splitTransformer = createSplitFunctionDeclarationsTransformer(split =>
-    collectedSplits.push(split)
-  );
-  const [sourceWithSplittedFunctions] = ts.transform(source, [
-    splitTransformer,
-  ]).transformed;
-
-  console.log(printer.printFile(sourceWithSplittedFunctions));
-  console.log(collectedSplits);
-  expect(1).toBe(1);
-});

test/inlineWrappedFunctions.test.ts

@@ -32,7 +32,7 @@ test('it can process nested calls of A2 with non identifiers as the first arg ',
       _VirtualDom_map_raw(fn, A2(styled.d9, add, context));
   `;
 
-  const { actual, expected } = transform(
+  const { actual, expected } = transformCode(
     initialCode,
     expectedOutputCode,
     createFunctionInlineTransformer()
@@ -70,7 +70,7 @@ test('it can process partial application inlining', () => {
   var res2 = func_raw(partialFunc2_a0, partialFunc2_a1, 3);
   `;
 
-  const { actual, expected } = transform(
+  const { actual, expected } = transformCode(
     initialCode,
     expectedOutputCode,
     createFunctionInlineTransformer()
@@ -100,7 +100,7 @@ test('it can inline functions that were wrapped by other functions', () => {
   var res = fullyApplied_raw(3, 4);
   `;
 
-  const { actual, expected } = transform(
+  const { actual, expected } = transformCode(
     initialCode,
     expectedOutputCode,
     createFunctionInlineTransformer()
@@ -135,7 +135,7 @@ test('it can inline functions that were wrapped by other functions even if they
     var res = fullyApplied_raw(3, 4);
   `;
 
-  const { actual, expected } = transform(
+  const { actual, expected } = transformCode(
     initialCode,
     expectedOutputCode,
     createFunctionInlineTransformer()
@@ -171,7 +171,7 @@ test('it can inline functions that were wrapped by other functions even if they
   var res = fullyApplied_raw(3, 4);
   `;
 
-  const { actual, expected } = transform(
+  const { actual, expected } = transformCode(
     initialCode,
     expectedOutputCode,
     createFunctionInlineTransformer()
@@ -180,7 +180,28 @@ test('it can inline functions that were wrapped by other functions even if they
   expect(actual).toBe(expected);
 });
 
-function transform(
+test('it can inline functions declared not via an identifier or lambda', () => {
+  const initialCode = `
+  var pow = F2(Math.pow);
+
+  var res = A2(pow, 2, 3);
+  `;
+
+  const expectedOutputCode = `
+  var pow_raw = Math.pow,  pow = F2(pow_raw);
+  
+  var res = pow_raw(2, 3);
+  `;
+
+  const { actual, expected } = transformCode(
+    initialCode,
+    expectedOutputCode,
+    createFunctionInlineTransformer()
+  );
+
+  expect(actual).toBe(expected);
+});
+export function transformCode(
   initialCode: string,
   expectedCode: string,
   transformer: ts.TransformerFactory<ts.SourceFile>

test/passUnwrappedFunctionsTransformer.test.ts

@@ -0,0 +1,170 @@
+import { transformCode } from './inlineWrappedFunctions.test';
+import { createPassUnwrappedFunctionsTransformer } from '../src/experiments/passUnwrappedFunctions';
+import { createInlineContext } from '../src/experiments/inlineWrappedFunctions';
+
+test('it can unwrap lambdas in place ', () => {
+  const initialCode = `
+    var f = function(func, a, b) {
+        return A2(func, a, b)
+    };
+
+    f(F2(function (a,b) {return a + b;}), 1, 2);
+`;
+
+  const expectedOutputCode = `
+    var f = function(func, a, b) {
+        return A2(func, a, b)
+    }, f_unwrapped = function(func, a, b) {
+        return func(a, b)
+    };
+
+    f_unwrapped(function (a,b) {return a + b;}, 1, 2);
+  `;
+
+  const { actual, expected } = transformCode(
+    initialCode,
+    expectedOutputCode,
+    createPassUnwrappedFunctionsTransformer(() => undefined)
+  );
+
+  expect(actual).toBe(expected);
+});
+
+test('it can deal with nesting too', () => {
+  const initialCode = `
+
+     var f = function(func, a, b) {
+        return A2(func, a, b)
+    };
+
+    var val = someFunc(
+          f(F2(function (a,b) {
+            return f(F2(function (c,d) {return  0} ), a, b)
+          }), 1,2)
+`;
+
+  const expectedOutputCode = `
+
+     var f = function(func, a, b) {
+        return A2(func, a, b)
+    }, f_unwrapped = function(func, a, b) {
+        return func(a, b)
+    };
+
+    var val = someFunc(
+          f_unwrapped(function (a,b) {
+            return f_unwrapped(function (c,d) {return  0}, a, b)
+          }, 1,2)
+  `;
+
+  const { actual, expected } = transformCode(
+    initialCode,
+    expectedOutputCode,
+    createPassUnwrappedFunctionsTransformer(() => undefined)
+  );
+
+  expect(actual).toBe(expected);
+});
+test('it can unwrap lambdas in place ', () => {
+  const initialCode = `
+    var f = function(func, a, b) {
+        return A2(func, a, b)
+    };
+
+    f(F2(function (a,b) {return a + b;}), 1, 2);
+`;
+
+  const expectedOutputCode = `
+    var f = function(func, a, b) {
+        return A2(func, a, b)
+    }, f_unwrapped = function(func, a, b) {
+        return func(a, b)
+    };
+
+    f_unwrapped(function (a,b) {return a + b;}, 1, 2);
+  `;
+
+  const { actual, expected } = transformCode(
+    initialCode,
+    expectedOutputCode,
+    createPassUnwrappedFunctionsTransformer(() => undefined)
+  );
+
+  expect(actual).toBe(expected);
+});
+
+test('it properly replaces names in recursion and all uses of the passed function', () => {
+  const initialCode = `
+    var g = function(func) {
+        return A2(func, 1, 1);
+    };
+
+    var f = function(func, a, b) {
+        return f(func, A2(func, a, a), A2(func, b, b)) + g(func);
+    };
+
+
+    f(F2(function (a,b) {return a + b;}), 1, 2);
+`;
+
+  const expectedOutputCode = `
+    var g = function(func) {
+        return A2(func, 1, 1);
+    }, g_unwrapped = function(func) {
+        return func(1, 1);
+    };
+
+    var f = function(func, a, b) {
+       return f(func, A2(func, a, a), A2(func, b, b)) + g(func);
+    }, f_unwrapped = function(func, a, b) {
+        return f_unwrapped(func, func(a, a), func(b, b)) + g_unwrapped(func);
+    };
+
+    f_unwrapped(function (a,b) {return a + b;}, 1, 2);
+  `;
+
+  const { actual, expected } = transformCode(
+    initialCode,
+    expectedOutputCode,
+    createPassUnwrappedFunctionsTransformer(() => undefined)
+  );
+
+  expect(actual).toBe(expected);
+});
+
+// commented out for now
+test('it can pass raw functions if there is one registered', () => {
+  const initialCode = `
+    var f = function(func, a, b) {
+        return A2(func, a, b)
+    };
+
+    f(wrappedFunc, 1, 2);
+`;
+
+  const expectedOutputCode = `
+    var f = function(func, a, b) {
+        return A2(func, a, b)
+    }, f_unwrapped = function(func, a, b) {
+        return func(a, b)
+    };
+
+    f_unwrapped(wrappedFunc_raw, 1, 2);
+  `;
+
+  const inlineContext = createInlineContext();
+
+  inlineContext.splits.set('wrappedFunc', {
+    arity: 2,
+    rawLambdaName: 'wrappedFunc_raw',
+    type: 'raw_func',
+  });
+
+  const { actual, expected } = transformCode(
+    initialCode,
+    expectedOutputCode,
+    createPassUnwrappedFunctionsTransformer(() => inlineContext)
+  );
+
+  expect(actual).toBe(expected);
+});