diff --git a/lib/transformer/injector_generator.dart b/lib/transformer/injector_generator.dart
index c863b33..923350c 100644
--- a/lib/transformer/injector_generator.dart
+++ b/lib/transformer/injector_generator.dart
@@ -364,15 +364,20 @@ class _Processor {
     var unit = lib.definingCompilationUnit.node;
     var transaction = resolver.createTextEditTransaction(lib);
 
+    FunctionDeclaration main = unit.declarations.firstWhere(
+        (d) => d is FunctionDeclaration && d.name.toString() == 'main',
+        orElse: () => null);
+
+    if (main == null) {
+      return;
+    }
+
     var imports = unit.directives.where((d) => d is ImportDirective);
     transaction.edit(imports.last.end, imports.last.end, '\nimport '
         "'${path.url.basenameWithoutExtension(id.path)}"
         "_generated_type_factory_maps.dart' show setStaticReflectorAsDefault;");
 
-    FunctionExpression main = unit.declarations.where((d) =>
-        d is FunctionDeclaration && d.name.toString() == 'main')
-        .first.functionExpression;
-    var body = main.body;
+    var body = main.functionExpression.body;
     if (body is BlockFunctionBody) {
       var location = body.beginToken.end;
       transaction.edit(location, location, '\n  setStaticReflectorAsDefault();');
diff --git a/test/transformer_test.dart b/test/transformer_test.dart
index 41592d4..ee85e70 100644
--- a/test/transformer_test.dart
+++ b/test/transformer_test.dart
@@ -769,6 +769,53 @@ main() {
             });
       });
 
+      it('transforms main with library', () {
+        return tests.applyTransformers(phases,
+            inputs: {
+              'a|web/main.dart': '''
+library main;
+import 'package:di/di.dart';
+
+main() {
+  print('abc');
+}''',
+              'a|web/lib.dart': '''
+library lib;
+part "a.dart";
+part "b.dart";''',
+              'a|web/a.dart': '''
+part of lib;
+Class A {
+}''',
+              'a|web/b.dart': '''
+part of lib;
+Class B{
+}'''
+            },
+            results: {
+              'a|web/main.dart': '''
+library main;
+import 'package:di/di.dart';
+import 'main_generated_type_factory_maps.dart' show setStaticReflectorAsDefault;
+
+main() {
+  setStaticReflectorAsDefault();
+  print('abc');
+}''',
+              'a|web/lib.dart': '''
+library lib;
+part "a.dart";
+part "b.dart";''',
+              'a|web/a.dart': '''
+part of lib;
+Class A {
+}''',
+              'a|web/b.dart': '''
+part of lib;
+Class B{
+}'''        });
+      });
+
       it('supports using a child of an injectable annotations as an injection marker', () {
         injectableAnnotations.add('di.annotations.Injectable');
         return generates(phases,